223 lines
8.4 KiB
Python
223 lines
8.4 KiB
Python
"""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"
|