Implement real-time WebSocket gateway for presence, DM, and world chat (US-4, US-2)

backend/tests/test_ws_gateway.py

@@ -0,0 +1,222 @@
+"""Integration tests for real-time WebSocket gateway:
+- Authentication (query param + first auth frame)
+- Rejection of unauthenticated clients
+- Presence on connect / disconnect
+- Subscriptions to world chat channels and DM conversations
+- Real-time event delivery from REST endpoints
+- Reconnect and ping/pong
+"""
+
+import asyncio
+import time
+import pytest
+from fastapi.testclient import TestClient
+from starlette.websockets import WebSocketDisconnect
+
+from app.main import app
+from app.services.presence import presence_store
+
+
+def _register(client: TestClient, username: str) -> dict:
+    resp = client.post(
+        "/api/auth/register",
+        json={
+            "username": username,
+            "email": f"{username}@example.com",
+            "password": "secret123",
+            "date_of_birth": "1990-01-01",
+        },
+    )
+    assert resp.status_code == 201, resp.text
+    token = resp.json()["access_token"]
+    headers = {"Authorization": f"Bearer {token}"}
+    me = client.get("/api/users/me", headers=headers)
+    assert me.status_code == 200, me.text
+    return {"headers": headers, "token": token, "id": me.json()["id"], "username": username}
+
+
+def _make_world(client: TestClient, headers: dict, visibility: str = "public") -> int:
+    r = client.post("/api/worlds/", headers=headers, json={"name": "WSWorld", "visibility": visibility})
+    assert r.status_code == 201, r.text
+    return r.json()["id"]
+
+
+def _make_friends(client: TestClient, user_a: dict, user_b: dict) -> None:
+    req = client.post("/api/friends/requests", json={"user_id": user_b["id"]}, headers=user_a["headers"])
+    assert req.status_code == 201, req.text
+    req_id = req.json()["id"]
+    acc = client.post(f"/api/friends/requests/{req_id}/accept", headers=user_b["headers"])
+    assert acc.status_code == 200, acc.text
+
+
+class TestWebSocketAuth:
+    def test_unauthenticated_connection_rejected(self):
+        client = TestClient(app)
+        # Invalid query param token rejected
+        with pytest.raises(WebSocketDisconnect) as exc:
+            with client.websocket_connect("/api/ws?token=invalid_token"):
+                pass
+        assert exc.value.code == 1008
+
+        # First frame with invalid token rejected
+        with pytest.raises(WebSocketDisconnect) as exc2:
+            with client.websocket_connect("/api/ws") as ws:
+                ws.send_json({"type": "auth", "token": "invalid_token"})
+                ws.receive_json()
+                ws.receive_json()
+        assert exc2.value.code == 1008
+
+    def test_auth_via_query_param(self):
+        client = TestClient(app)
+        alice = _register(client, "ws_auth_alice")
+
+        with client.websocket_connect(f"/api/ws?token={alice['token']}") as ws:
+            msg = ws.receive_json()
+            assert msg["type"] == "auth_ok"
+            assert msg["user_id"] == alice["id"]
+            assert msg["username"] == "ws_auth_alice"
+
+    def test_auth_via_first_message(self):
+        client = TestClient(app)
+        bob = _register(client, "ws_auth_bob")
+
+        with client.websocket_connect("/api/ws") as ws:
+            ws.send_json({"type": "auth", "token": bob["token"]})
+            msg = ws.receive_json()
+            assert msg["type"] == "auth_ok"
+            assert msg["user_id"] == bob["id"]
+            assert msg["username"] == "ws_auth_bob"
+
+
+class TestWebSocketPresence:
+    def test_presence_on_connect_and_disconnect(self):
+        import threading
+        client = TestClient(app)
+        alice = _register(client, "ws_pres_alice")
+        bob = _register(client, "ws_pres_bob")
+        _make_friends(client, alice, bob)
+
+        # Initially both offline
+        assert not asyncio.run(presence_store.is_online(alice["id"]))
+
+        # Bob connects and listens for presence events
+        bob_events = []
+
+        def run_alice():
+            time.sleep(0.2)
+            c_alice = TestClient(app)
+            with c_alice.websocket_connect(f"/api/ws?token={alice['token']}") as alice_ws:
+                assert alice_ws.receive_json()["type"] == "auth_ok"
+                time.sleep(0.3)
+                alice_ws.close()
+                time.sleep(0.3)
+
+        with client.websocket_connect(f"/api/ws?token={bob['token']}") as bob_ws:
+            assert bob_ws.receive_json()["type"] == "auth_ok"
+
+            alice_thread = threading.Thread(target=run_alice)
+            alice_thread.start()
+
+            # Bob receives Alice's online event
+            online_event = bob_ws.receive_json()
+            assert online_event["type"] == "presence"
+            assert online_event["user_id"] == alice["id"]
+            assert online_event["is_online"] is True
+
+            alice_thread.join()
+
+            # Bob receives Alice's offline event
+            offline_event = bob_ws.receive_json()
+            assert offline_event["type"] == "presence"
+            assert offline_event["user_id"] == alice["id"]
+            assert offline_event["is_online"] is False
+
+        assert not asyncio.run(presence_store.is_online(alice["id"]))
+
+
+class TestWebSocketChatDelivery:
+    def test_world_chat_subscription_and_delivery(self):
+        client = TestClient(app)
+        alice = _register(client, "ws_chat_alice")
+        bob = _register(client, "ws_chat_bob")
+        world_id = _make_world(client, alice["headers"])
+
+        with client.websocket_connect(f"/api/ws?token={bob['token']}") as bob_ws:
+            assert bob_ws.receive_json()["type"] == "auth_ok"
+
+            # Bob subscribes to world
+            bob_ws.send_json({"type": "subscribe_world", "world_id": world_id})
+            sub_ack = bob_ws.receive_json()
+            assert sub_ack["type"] == "subscribed"
+            assert sub_ack["channel"] == "world"
+            assert sub_ack["id"] == world_id
+
+            # Alice posts message to world via REST
+            rest_resp = client.post(
+                f"/api/chat/world/{world_id}/messages",
+                headers=alice["headers"],
+                json={"body": "Hello real-time world!"},
+            )
+            assert rest_resp.status_code == 201, rest_resp.text
+            msg_id = rest_resp.json()["id"]
+
+            # Bob receives message in real time over WebSocket
+            ws_msg = bob_ws.receive_json()
+            assert ws_msg["type"] == "world_chat"
+            assert ws_msg["world_id"] == world_id
+            assert ws_msg["message"]["id"] == msg_id
+            assert ws_msg["message"]["body"] == "Hello real-time world!"
+            assert ws_msg["message"]["author"]["username"] == "ws_chat_alice"
+
+    def test_dm_chat_subscription_and_delivery(self):
+        client = TestClient(app)
+        alice = _register(client, "ws_dm_alice")
+        bob = _register(client, "ws_dm_bob")
+
+        # Initial DM to create conversation
+        init_dm = client.post(
+            f"/api/chat/dm/{bob['id']}/messages",
+            headers=alice["headers"],
+            json={"body": "First DM to initiate"},
+        )
+        assert init_dm.status_code == 201, init_dm.text
+        conv_id = init_dm.json()["conversation_id"]
+
+        with client.websocket_connect(f"/api/ws?token={bob['token']}") as bob_ws:
+            assert bob_ws.receive_json()["type"] == "auth_ok"
+
+            # Bob subscribes to this conversation
+            bob_ws.send_json({"type": "subscribe_dm", "conversation_id": conv_id})
+            sub_ack = bob_ws.receive_json()
+            assert sub_ack["type"] == "subscribed"
+            assert sub_ack["channel"] == "dm"
+            assert sub_ack["id"] == conv_id
+
+            # Alice sends another DM via REST
+            rest_dm = client.post(
+                f"/api/chat/dm/{bob['id']}/messages",
+                headers=alice["headers"],
+                json={"body": "Real-time DM ping!"},
+            )
+            assert rest_dm.status_code == 201, rest_dm.text
+
+            # Bob receives real-time DM over WebSocket
+            ws_dm = bob_ws.receive_json()
+            assert ws_dm["type"] == "dm_chat"
+            assert ws_dm["conversation_id"] == conv_id
+            assert ws_dm["message"]["body"] == "Real-time DM ping!"
+
+
+class TestWebSocketPingAndHeartbeat:
+    def test_ping_pong_and_heartbeat(self):
+        client = TestClient(app)
+        alice = _register(client, "ws_ping_alice")
+
+        with client.websocket_connect(f"/api/ws?token={alice['token']}") as ws:
+            assert ws.receive_json()["type"] == "auth_ok"
+
+            ws.send_json({"type": "ping"})
+            assert ws.receive_json()["type"] == "pong"
+
+            ws.send_json({"type": "heartbeat"})
+            assert ws.receive_json()["type"] == "heartbeat_ack"