Update tests

chat/src/event_config.py

@@ -92,20 +92,6 @@ def _get_temperature(self):
     def _get_text_key(self):
         return self._get_payload_value_with_superuser_check("text_key", TEXT_KEY)
 
-    def debug_message(self):
-        return {
-            "type": "debug",
-            "message": {
-                "k": self.k,
-                "prompt": self.prompt_text,
-                "question": self.question,
-                "ref": self.ref,
-                "size": self.ref,
-                "temperature": self.temperature,
-                "text_key": self.text_key,
-            },
-        }
-
     def setup_websocket(self, socket=None):
         if socket is None:
             connection_id = self.request_context.get("connectionId")
@@ -120,11 +106,3 @@ def setup_websocket(self, socket=None):
     def _is_debug_mode_enabled(self):
         debug = self.payload.get("debug", False)
         return debug and self.api_token.is_superuser()
-
-    def _to_bool(self, val):
-        """Converts a value to boolean. If the value is a string, it considers
-        "", "no", "false", "0" as False. Otherwise, it returns the boolean of the value.
-        """
-        if isinstance(val, str):
-            return val.lower() not in ["", "no", "false", "0"]
-        return bool(val)

chat/src/websocket.py

@@ -8,9 +8,9 @@ def __init__(self, client=None, endpoint_url=None, connection_id=None, ref=None)
         self.ref = ref if ref else {}
 
     def send(self, data):
-        # if isinstance(data, str):
-        #     data = {"message": data}
-        # data["ref"] = self.ref
+        if isinstance(data, str):
+            data = {"message": data}
+        data["ref"] = self.ref
         data_as_bytes = bytes(json.dumps(data), "utf-8")
 
         if self.connection_id == "debug":

chat/test/handlers/test_chat.py

@@ -9,8 +9,8 @@
 from unittest import mock, TestCase
 from unittest.mock import patch
 from handlers.chat import handler
+from agent.search_agent import SearchAgent
 from helpers.apitoken import ApiToken
-from helpers.response import Response
 from websocket import Websocket
 from event_config import EventConfig
 
@@ -55,7 +55,7 @@ def mock_response(**kwargs):
         "AZURE_OPENAI_RESOURCE_NAME": "test",
     },
 )
-@mock.patch.object(Response, "prepare_response", lambda _: mock_response())
+
 class TestHandler(TestCase):
     def test_handler_unauthorized(self):        
         event = {"socket": Websocket(client=MockClient(), endpoint_url="test", connection_id="test", ref="test")}
@@ -66,36 +66,6 @@ def test_handler_success(self, mock_is_logged_in):
       mock_is_logged_in.return_value = True
       event = {"socket": Websocket(client=MockClient(), endpoint_url="test", connection_id="test", ref="test"), "body": '{"question": "Question?"}' }
       self.assertEqual(handler(event, MockContext()), {'statusCode': 200})
-
-    @patch.object(ApiToken, 'is_logged_in')
-    @patch.object(ApiToken, 'is_superuser')
-    @patch.object(EventConfig, '_is_debug_mode_enabled')
-    def test_handler_debug_mode(self, mock_is_debug_enabled, mock_is_superuser, mock_is_logged_in):
-      mock_is_debug_enabled.return_value = True
-      mock_is_logged_in.return_value = True
-      mock_is_superuser.return_value = True
-      mock_client = MockClient()
-      mock_websocket = Websocket(client=mock_client, endpoint_url="test", connection_id="test", ref="test")
-      event = {"socket": mock_websocket, "debug": True, "body": '{"question": "Question?"}' }
-      handler(event, MockContext())
-      response = json.loads(mock_client.received_data)
-      expected_keys = {"attributes", "azure_endpoint", "deployment_name"}
-      received_keys = response.keys()
-      self.assertTrue(expected_keys.issubset(received_keys))
-
-    @patch.object(ApiToken, 'is_logged_in')
-    @patch.object(ApiToken, 'is_superuser')
-    def test_handler_debug_mode_for_superusers_only(self, mock_is_superuser, mock_is_logged_in):
-      mock_is_logged_in.return_value = True
-      mock_is_superuser.return_value = False
-      mock_client = MockClient()
-      mock_websocket = Websocket(client=mock_client, endpoint_url="test", connection_id="test", ref="test")
-      event = {"socket": mock_websocket, "body": '{"question": "Question?", "debug": "true"}'}
-      handler(event, MockContext())
-      response = json.loads(mock_client.received_data)
-      expected_keys = {"answer", "ref"}
-      received_keys = set(response.keys())
-      self.assertSetEqual(received_keys, expected_keys)
 
     @patch.object(ApiToken, 'is_logged_in')
     def test_handler_question_missing(self, mock_is_logged_in):

chat/test/test_event_config.py

@@ -7,89 +7,35 @@
 from event_config import EventConfig
 from unittest import TestCase, mock
 
-
-class TestEventConfigWithoutAzureResource(TestCase):
-    def test_requires_an_azure_resource(self):
-        with self.assertRaises(EnvironmentError):
-            EventConfig()
-
-
-@mock.patch.dict(
-    os.environ,
-    {
-        "AZURE_OPENAI_RESOURCE_NAME": "test",
-    },
-)
 class TestEventConfig(TestCase):
-    def test_fetches_attributes_from_vector_database(self):
-        os.environ.pop("AZURE_OPENAI_RESOURCE_NAME", None)
-        with self.assertRaises(EnvironmentError):
-            EventConfig()
-
     def test_defaults(self):
-        actual = EventConfig(event={"body": json.dumps({"attributes": ["title"]})})
-        expected_defaults = {"azure_endpoint": "https://test.openai.azure.com/"}
-        self.assertEqual(actual.azure_endpoint, expected_defaults["azure_endpoint"])
+        actual = EventConfig(event={"body": json.dumps({"question": "Question?"})})
+        self.assertEqual(actual.question, "Question?")
 
     def test_attempt_override_without_superuser_status(self):
         actual = EventConfig(
             event={
                 "body": json.dumps(
                     {
-                        "azure_resource_name": "new_name_for_test",
-                        "attributes": ["title", "subject", "date_created"],
-                        "index": "testIndex",
                         "k": 100,
-                        "openai_api_version": "2024-01-01",
                         "question": "test question",
                         "ref": "test ref",
                         "size": 90,
                         "temperature": 0.9,
-                        "text_key": "accession_number",
                     }
                 )
             }
         )
         expected_output = {
-            "attributes": EventConfig.DEFAULT_ATTRIBUTES,
-            "azure_endpoint": "https://test.openai.azure.com/",
             "k": 40,
-            "openai_api_version": "2024-02-01",
             "question": "test question",
-            "size": 20,
             "ref": "test ref",
+            "size": 20,
             "temperature": 0.2,
-            "text_key": "id",
         }
-        self.assertEqual(actual.azure_endpoint, expected_output["azure_endpoint"])
-        self.assertEqual(actual.attributes, expected_output["attributes"])
         self.assertEqual(actual.k, expected_output["k"])
-        self.assertEqual(
-            actual.openai_api_version, expected_output["openai_api_version"]
-        )
         self.assertEqual(actual.question, expected_output["question"])
         self.assertEqual(actual.ref, expected_output["ref"])
+        self.assertEqual(actual.size, expected_output["size"])
         self.assertEqual(actual.temperature, expected_output["temperature"])
-        self.assertEqual(actual.text_key, expected_output["text_key"])
-
-    def test_debug_message(self):
-        self.assertEqual(
-            EventConfig(
-                event={"body": json.dumps({"attributes": ["source"]})}
-            ).debug_message()["type"],
-            "debug",
-        )
-
-    def test_to_bool(self):
-        self.assertEqual(EventConfig(event={"body": json.dumps({"attributes": ["source"]})})._to_bool(""), False)
-        self.assertEqual(EventConfig(event={"body": json.dumps({"attributes": ["source"]})})._to_bool("0"), False)
-        self.assertEqual(EventConfig(event={"body": json.dumps({"attributes": ["source"]})})._to_bool("no"), False)
-        self.assertEqual(EventConfig(event={"body": json.dumps({"attributes": ["source"]})})._to_bool("false"), False)
-        self.assertEqual(EventConfig(event={"body": json.dumps({"attributes": ["source"]})})._to_bool("False"), False)
-        self.assertEqual(EventConfig(event={"body": json.dumps({"attributes": ["source"]})})._to_bool("FALSE"), False)
-        self.assertEqual(EventConfig(event={"body": json.dumps({"attributes": ["source"]})})._to_bool("no"), False)
-        self.assertEqual(EventConfig(event={"body": json.dumps({"attributes": ["source"]})})._to_bool("No"), False)
-        self.assertEqual(EventConfig(event={"body": json.dumps({"attributes": ["source"]})})._to_bool("NO"), False)
-        self.assertEqual(EventConfig(event={"body": json.dumps({"attributes": ["source"]})})._to_bool("true"), True)
-        self.assertEqual(EventConfig(event={"body": json.dumps({"attributes": ["source"]})})._to_bool(True), True)
-        self.assertEqual(EventConfig(event={"body": json.dumps({"attributes": ["source"]})})._to_bool(False), False)
+