test: Fix websocket tests

deeppavlov · Sep 17, 2024 · 20a0c0e · 20a0c0e
1 parent 073a5db
commit 20a0c0e
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 17 deletions.
diff --git a/backend/chatsky_ui/tests/api/test_bot.py b/backend/chatsky_ui/tests/api/test_bot.py
@@ -156,15 +156,15 @@ async def test_get_run_logs(mocker, pagination):
 async def test_connect(mocker):
     websocket = mocker.AsyncMock()
     websocket_manager = mocker.AsyncMock()
-    websocket_manager.disconnect = mocker.MagicMock()
-    run_manager = mocker.MagicMock()
-    run_process = mocker.MagicMock()
+    websocket_manager.disconnect = mocker.AsyncMock()
+    run_manager = mocker.AsyncMock()
+    run_process = mocker.AsyncMock()
     run_manager.processes = {RUN_ID: run_process}
+    run_manager.get_status = mocker.AsyncMock(return_value=Status.ALIVE)
     mocker.patch.object(websocket, "query_params", {"run_id": str(RUN_ID)})
 
     await connect(websocket, websocket_manager, run_manager)
 
-    websocket_manager.connect.assert_awaited_once_with(websocket)
+    websocket_manager.connect.assert_awaited_once_with(RUN_ID, websocket)
     websocket_manager.send_process_output_to_websocket.assert_awaited_once_with(RUN_ID, run_manager, websocket)
     websocket_manager.forward_websocket_messages_to_process.assert_awaited_once_with(RUN_ID, run_manager, websocket)
-    websocket_manager.disconnect.assert_called_once_with(websocket)
diff --git a/backend/chatsky_ui/tests/services/test_websocket_manager.py b/backend/chatsky_ui/tests/services/test_websocket_manager.py
@@ -2,56 +2,64 @@
 from fastapi import WebSocket
 
 
+RUN_ID = 42
+
+
 class TestWebSocketManager:
     @pytest.mark.asyncio
     async def test_connect(self, mocker, websocket_manager):
-        mocked_websocket = mocker.MagicMock(spec=WebSocket)
+        mocked_websocket = mocker.MagicMock()
+        mocked_websocket.accept = mocker.AsyncMock()
 
-        await websocket_manager.connect(mocked_websocket)
+        await websocket_manager.connect(RUN_ID, mocked_websocket)
 
         mocked_websocket.accept.assert_awaited_once_with()
-        assert mocked_websocket in websocket_manager.active_connections
+        assert mocked_websocket == websocket_manager.active_connections[RUN_ID]["websocket"]
 
     @pytest.mark.asyncio
     async def test_disconnect(self, mocker, websocket_manager):
         mocked_websocket = mocker.MagicMock(spec=WebSocket)
-        websocket_manager.active_connections.append(mocked_websocket)
+        websocket_manager.active_connections[RUN_ID] = {"websocket": mocked_websocket, "chat": {}}
         websocket_manager.pending_tasks[mocked_websocket] = set()
 
-        websocket_manager.disconnect(mocked_websocket)
+        await websocket_manager.disconnect(RUN_ID, mocked_websocket)
 
         assert mocked_websocket not in websocket_manager.pending_tasks
-        assert mocked_websocket not in websocket_manager.active_connections
+        assert RUN_ID not in websocket_manager.active_connections
 
     @pytest.mark.asyncio
     async def test_send_process_output_to_websocket(self, mocker, websocket_manager):
-        run_id = 42
         awaited_response = "Hello from DF-Designer"
 
+        mocked_websocket = mocker.MagicMock(spec=WebSocket)
+        websocket_manager.active_connections[RUN_ID] = {"websocket": mocked_websocket, "chat": {"messages": []}}
+
         websocket = mocker.AsyncMock()
         run_manager = mocker.MagicMock()
         run_process = mocker.MagicMock()
         run_process.read_stdout = mocker.AsyncMock(side_effect=[awaited_response.encode(), None])
-        run_manager.processes = {run_id: run_process}
+        run_manager.processes = {RUN_ID: run_process}
 
-        await websocket_manager.send_process_output_to_websocket(run_id, run_manager, websocket)
+        await websocket_manager.send_process_output_to_websocket(RUN_ID, run_manager, websocket)
 
         assert run_process.read_stdout.call_count == 2
         websocket.send_text.assert_awaited_once_with(awaited_response)
 
     @pytest.mark.asyncio
     async def test_forward_websocket_messages_to_process(self, mocker, websocket_manager):
-        run_id = 42
         awaited_message = "Hello from DF-Designer"
 
+        mocked_websocket = mocker.MagicMock(spec=WebSocket)
+        websocket_manager.active_connections[RUN_ID] = {"websocket": mocked_websocket, "chat": {"messages": []}}
+
         websocket = mocker.AsyncMock()
         websocket.receive_text = mocker.AsyncMock(side_effect=[awaited_message, None])
         run_manager = mocker.MagicMock()
         run_process = mocker.MagicMock()
         run_process.write_stdin = mocker.AsyncMock()
-        run_manager.processes = {run_id: run_process}
+        run_manager.processes = {RUN_ID: run_process}
 
-        await websocket_manager.forward_websocket_messages_to_process(run_id, run_manager, websocket)
+        await websocket_manager.forward_websocket_messages_to_process(RUN_ID, run_manager, websocket)
 
         assert websocket.receive_text.await_count == 2
         run_process.write_stdin.assert_called_once_with(awaited_message.encode() + b"\n")