Fix Github action tests

nulib · Feb 1, 2024 · b4bc082 · b4bc082
1 parent 61f1378
commit b4bc082
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/chat/test/handlers/test_streaming_socket_callback_handler.py b/chat/test/handlers/test_streaming_socket_callback_handler.py
@@ -4,14 +4,17 @@
 )
 from src.websocket import Websocket
 
+class MockClient:
+    def post_to_connection(self, Data, ConnectionId):
+        return Data
 
 class TestMyStreamingSocketCallbackHandler(TestCase):
     def test_on_new_llm_token(self):
-        handler = StreamingSocketCallbackHandler(Websocket(), False)
+        handler = StreamingSocketCallbackHandler(Websocket(client=MockClient()), False)
         result = handler.on_llm_new_token(token="test")
         self.assertEqual(result, "test")
         self.assertFalse(handler.debug_mode)
 
     def test_debug_mode(self):
-        handler = StreamingSocketCallbackHandler(Websocket(), debug_mode=True)
+        handler = StreamingSocketCallbackHandler(Websocket(client=MockClient()), debug_mode=True)
         self.assertTrue(handler.debug_mode)