Get chat and sync chat tests passing

chat/test/handlers/test_chat.py

@@ -8,7 +8,6 @@
 sys.path.append('./src')
 
 from handlers.chat import handler
-from handlers.model import set_model_override
 from helpers.apitoken import ApiToken
 from langchain_core.language_models.fake_chat_models import FakeListChatModel
 from langgraph.checkpoint.memory import MemorySaver
@@ -37,8 +36,8 @@ def test_handler_unauthorized(self, mock_create_saver, mock_is_logged_in):
 
     @patch.object(ApiToken, 'is_logged_in', return_value=True)
     @patch('agent.search_agent.checkpoint_saver', return_value=MemorySaver())
-    def test_handler_success(self, mock_create_saver, mock_is_logged_in):
-        set_model_override(FakeListChatModel(responses=["fake response"]))
+    @patch('handlers.chat.chat_model', return_value=FakeListChatModel(responses=["fake response"]))
+    def test_handler_success(self, mock_chat_model, mock_create_saver, mock_is_logged_in):
         event = {
             "socket": Websocket(client=MockClient(), endpoint_url="test", connection_id="test", ref="test"),
             "body": '{"question": "Question?"}'

chat/test/handlers/test_chat_sync.py

@@ -1,11 +1,14 @@
 # ruff: noqa: E402
 
+import json
 import os
 import sys
+from langchain_core.language_models.fake_chat_models import FakeListChatModel
+from langgraph.checkpoint.memory import MemorySaver
 
 sys.path.append('./src')
 
-from unittest import mock, TestCase
+from unittest import TestCase
 from unittest.mock import patch
 from handlers.chat_sync import handler
 from helpers.apitoken import ApiToken
@@ -14,22 +17,31 @@ class MockContext:
    def __init__(self):
       self.log_stream_name = 'test'
 
-@mock.patch.dict(
-    os.environ,
-    {
-        "AZURE_OPENAI_RESOURCE_NAME": "test",
-    },
-)
 class TestHandler(TestCase):
     def test_handler_unauthorized(self):        
         self.assertEqual(handler({"body": '{ "question": "Question?"}'}, MockContext()), {'body': 'Unauthorized', 'statusCode': 401})
 
-    @patch.object(ApiToken, 'is_logged_in')
+    @patch.object(ApiToken, 'is_logged_in', return_value = True)
     def test_no_question(self, mock_is_logged_in):
-      mock_is_logged_in.return_value = True
       self.assertEqual(handler({"body": '{ "question": ""}'}, MockContext()), {'statusCode': 400, 'body': 'Question cannot be blank'})
 
-    @patch.object(ApiToken, 'is_logged_in')
-    def test_handler_success(self, mock_is_logged_in):
-      mock_is_logged_in.return_value = True
-      self.assertEqual(handler({"body": '{"question": "Question?"}'}, MockContext()), {'statusCode': 200})
+    @patch.object(ApiToken, 'is_logged_in', return_value = True)
+    @patch("agent.search_agent.checkpoint_saver", return_value=MemorySaver())
+    @patch('handlers.chat_sync.chat_model', return_value=FakeListChatModel(responses=["fake response"]))
+    def test_handler_success(self, mock_chat_model, mock_create_saver, mock_is_logged_in):
+        expected_body = {
+            "answer": ["fake response"], 
+            "is_dev_team": False, 
+            "is_superuser": False, 
+            "k": 40, 
+            "model": "us.anthropic.claude-3-5-sonnet-20241022-v2:0", 
+            "question": "Question?", 
+            "ref": "test_ref", 
+            "artifacts": [], 
+            "token_counts": {}
+        }
+        response = handler({"body": '{"question": "Question?", "ref": "test_ref"}'}, MockContext())
+
+        self.assertEqual(json.loads(response.get("body")), expected_body)
+        self.assertEqual(response.get("statusCode"), 200)
+        self.assertEqual(response.get("headers", {}).get("Content-Type"), "application/json")