improve ws tests

cheshire-cat-ai · Feb 8, 2025 · 937cb01 · 937cb01
1 parent 5950fa7
commit 937cb01
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 12 deletions.
diff --git a/core/tests/routes/test_websocket.py b/core/tests/routes/test_websocket.py
@@ -85,22 +85,23 @@ def test_websocket_multiple_connections(client):
 
         with client.websocket_connect("/ws/Caterpillar") as websocket2:
 
-                # send ws message
-                websocket2.send_json(mex)
-                # get reply
-                reply2 = websocket2.receive_json()
+            # send ws message
+            websocket2.send_json(mex)
+            # get reply
+            reply2 = websocket2.receive_json()
 
-                # two connections open
-                assert set(client.app.state.websocket_manager.connections.keys()) \
-                                == {"Alice", "Caterpillar"}
+            # two connections open
+            ws_users = client.app.state.websocket_manager.connections.keys()
+            assert set(ws_users) == {"Alice", "Caterpillar"}
 
-                websocket2.close()
+        # one connection open
+        time.sleep(0.5)
+        ws_users = client.app.state.websocket_manager.connections.keys()
+        assert set(ws_users) == {"Alice"}
 
         # get reply
         reply = websocket.receive_json()
 
-        websocket.close()
-
     check_correct_websocket_reply(reply)
     check_correct_websocket_reply(reply2)
 

diff --git a/core/tests/utils.py b/core/tests/utils.py
@@ -15,8 +15,6 @@ def send_websocket_message(msg, client, user_id="user", query_params=None):
         # get reply
         reply = websocket.receive_json()
 
-        websocket.close()
-
     return reply