"""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"