"""WebSocket live-play tests via kaya's ASGI websocket transport.""" from __future__ import annotations import json import unittest from typing import Any, Dict, Mapping, Sequence from httpx import ASGITransport, AsyncClient from httpx_ws import WebSocketDisconnect as HttpxWebSocketDisconnect, aconnect_ws from httpx_ws.transport import ASGIWebSocketTransport from kaya.core import WebSocket, WebSocketDisconnect, WebSocketMessage from helpers import async_test, make_platform, make_user, oidc_user, ws_users from tavolo.platform import GameSession, Seat class AbruptCloseSocket(WebSocket): """A socket whose peer vanished without a close handshake. kaya surfaces this from the transport: ``receive()`` reports a ``close`` message so iteration ends normally, while sends raise ``WebSocketDisconnect``. """ def __init__(self, path: str, gone: bool = False) -> None: self.path = path self.query_string = "" self.scheme = "ws" self.headers: Mapping[str, Sequence[str]] = {} self.client = ("127.0.0.1", 12345) self.server = ("127.0.0.1", 80) self.sent: list[Dict[str, Any]] = [] self.close_codes: list[int] = [] self.gone = gone async def accept(self, headers: Any = None) -> None: pass async def receive(self) -> WebSocketMessage: self.gone = True return WebSocketMessage(kind="close") async def send_text(self, data: str) -> None: if self.gone: raise WebSocketDisconnect(code=1006) self.sent.append(json.loads(data)) async def send_bytes(self, data: bytes) -> None: if self.gone: raise WebSocketDisconnect(code=1006) async def close(self, code: int = 1000) -> None: self.close_codes.append(code) async def __anext__(self) -> WebSocketMessage: await self.receive() raise StopAsyncIteration class RecordingSocket(WebSocket): """A socket that only records pre-accept rejections.""" def __init__(self, path: str) -> None: self.path = path self.query_string = "" self.scheme = "ws" self.headers: Mapping[str, Sequence[str]] = {} self.client = ("127.0.0.1", 12345) self.server = ("127.0.0.1", 80) self.close_codes: list[int] = [] async def accept(self, headers: Any = None) -> None: raise AssertionError("a rejected socket must not be accepted") async def receive(self) -> WebSocketMessage: raise AssertionError("a rejected socket must not receive") async def send_text(self, data: str) -> None: raise AssertionError("a rejected socket must not send") async def send_bytes(self, data: bytes) -> None: raise AssertionError("a rejected socket must not send") async def close(self, code: int = 1000) -> None: self.close_codes.append(code) async def __anext__(self) -> WebSocketMessage: raise AssertionError("a rejected socket must not receive") class WebSocketTest(unittest.TestCase): async def _started_game(self, client: AsyncClient, oidc, target: int = 3) -> dict: """Create a game and seat both players; return the started state.""" with oidc_user(oidc, "alice"): created = await client.post( "/api/games", json={"options": {"target": target}} ) code = created.json()["join_code"] with oidc_user(oidc, "bob"): response = await client.post("/api/games/join", json={"code": code}) return response.json() @async_test async def test_move_updates_all_connections(self) -> None: app, platform, _ = make_platform() api_transport = ASGITransport(app=app) async with AsyncClient(transport=api_transport, base_url="http://127.0.0.1") as client: state = await self._started_game(client, platform.oidc) game_id = state["id"] ws_transport = ASGIWebSocketTransport(app=app) async with AsyncClient(transport=ws_transport, base_url="http://testserver") as ws_client: with ws_users([make_user("bob"), make_user("alice")]): async with aconnect_ws(f"/ws/games/{game_id}", ws_client) as bob_ws: first = await bob_ws.receive_json() self.assertEqual("state", first["type"]) self.assertEqual("dummy", first["game"]["game_type"]) self.assertEqual(0, first["game"]["plays"]) self.assertEqual(game_id, first["game"]["id"]) async with aconnect_ws(f"/ws/games/{game_id}", ws_client) as alice_ws: alice_first = await alice_ws.receive_json() self.assertEqual("state", alice_first["type"]) await bob_ws.send_json({"action": "play"}) bob_update = await bob_ws.receive_json() alice_update = await alice_ws.receive_json() for update in (bob_update, alice_update): self.assertEqual("state", update["type"]) self.assertEqual(1, update["game"]["plays"]) @async_test async def test_unknown_action_returns_error(self) -> None: app, platform, _ = make_platform() api_transport = ASGITransport(app=app) async with AsyncClient(transport=api_transport, base_url="http://127.0.0.1") as client: state = await self._started_game(client, platform.oidc) game_id = state["id"] ws_transport = ASGIWebSocketTransport(app=app) async with AsyncClient(transport=ws_transport, base_url="http://testserver") as ws_client: with ws_users([make_user("bob")]): async with aconnect_ws(f"/ws/games/{game_id}", ws_client) as bob_ws: await bob_ws.receive_json() await bob_ws.send_json({"action": "dance"}) error = await bob_ws.receive_json() self.assertEqual("error", error["type"]) self.assertEqual("illegal_move", error["code"]) @async_test async def test_state_action_resyncs(self) -> None: app, platform, _ = make_platform() api_transport = ASGITransport(app=app) async with AsyncClient(transport=api_transport, base_url="http://127.0.0.1") as client: state = await self._started_game(client, platform.oidc) game_id = state["id"] ws_transport = ASGIWebSocketTransport(app=app) async with AsyncClient(transport=ws_transport, base_url="http://testserver") as ws_client: with ws_users([make_user("bob")]): async with aconnect_ws(f"/ws/games/{game_id}", ws_client) as bob_ws: await bob_ws.receive_json() await bob_ws.send_json({"action": "state"}) resent = await bob_ws.receive_json() self.assertEqual("state", resent["type"]) @async_test async def test_abrupt_disconnect_is_not_an_error(self) -> None: app, platform, _ = make_platform() api_transport = ASGITransport(app=app) async with AsyncClient(transport=api_transport, base_url="http://127.0.0.1") as client: state = await self._started_game(client, platform.oidc) game_id = state["id"] ws = AbruptCloseSocket(f"/ws/games/{game_id}") with ws_users([make_user("bob")]): # kaya reports a peer that vanished without a close handshake as # a close message; the handler must exit cleanly. await app.handle_websocket(ws) self.assertEqual("state", ws.sent[0]["type"]) self.assertEqual([], ws.close_codes) @async_test async def test_send_on_vanished_peer_is_ignored(self) -> None: app, platform, _ = make_platform() api_transport = ASGITransport(app=app) async with AsyncClient(transport=api_transport, base_url="http://127.0.0.1") as client: state = await self._started_game(client, platform.oidc) game_id = state["id"] ws = AbruptCloseSocket(f"/ws/games/{game_id}", gone=True) with ws_users([make_user("bob")]): # Sending raises WebSocketDisconnect; the handler swallows it # rather than letting it escape the forward task. await app.handle_websocket(ws) self.assertEqual([], ws.sent) self.assertEqual([], ws.close_codes) @async_test async def test_game_over_broadcast(self) -> None: app, platform, _ = make_platform() api_transport = ASGITransport(app=app) async with AsyncClient(transport=api_transport, base_url="http://127.0.0.1") as client: # target 1: the first play ends the match. state = await self._started_game(client, platform.oidc, target=1) game_id = state["id"] ws_transport = ASGIWebSocketTransport(app=app) async with AsyncClient(transport=ws_transport, base_url="http://testserver") as ws_client: with ws_users([make_user("alice"), make_user("bob")]): async with aconnect_ws(f"/ws/games/{game_id}", ws_client) as alice_ws: await alice_ws.receive_json() async with aconnect_ws(f"/ws/games/{game_id}", ws_client) as bob_ws: await bob_ws.receive_json() await bob_ws.send_json({"action": "play"}) # Both connections see the final state... alice_final = await alice_ws.receive_json() bob_final = await bob_ws.receive_json() self.assertTrue(alice_final["game"]["finished"]) self.assertTrue(bob_final["game"]["finished"]) # ...followed by the game_over announcement. alice_over = await alice_ws.receive_json() bob_over = await bob_ws.receive_json() self.assertEqual("game_over", alice_over["type"]) self.assertEqual("game_over", bob_over["type"]) @async_test async def test_unknown_game_is_closed(self) -> None: app, _, _ = make_platform() ws_transport = ASGIWebSocketTransport(app=app) async with AsyncClient(transport=ws_transport, base_url="http://testserver") as ws_client: with ws_users([make_user("alice")]): with self.assertRaises(HttpxWebSocketDisconnect) as caught: async with aconnect_ws("/ws/games/no-such-game", ws_client): pass self.assertEqual(404, caught.exception.code) @async_test async def test_non_player_is_closed(self) -> None: app, platform, _ = make_platform() api_transport = ASGITransport(app=app) async with AsyncClient(transport=api_transport, base_url="http://127.0.0.1") as client: state = await self._started_game(client, platform.oidc) game_id = state["id"] ws_transport = ASGIWebSocketTransport(app=app) async with AsyncClient(transport=ws_transport, base_url="http://testserver") as ws_client: with ws_users([make_user("mallory")]): with self.assertRaises(HttpxWebSocketDisconnect) as caught: async with aconnect_ws(f"/ws/games/{game_id}", ws_client): pass self.assertEqual(403, caught.exception.code) @async_test async def test_unauthenticated_is_closed(self) -> None: app, _, _ = make_platform() ws_transport = ASGIWebSocketTransport(app=app) async with AsyncClient(transport=ws_transport, base_url="http://testserver") as ws_client: with ws_users([]): with self.assertRaises(HttpxWebSocketDisconnect) as caught: async with aconnect_ws("/ws/games/whatever", ws_client): pass self.assertEqual(401, caught.exception.code) class RejectionStatusTest(unittest.TestCase): """Pre-accept rejections must carry HTTP statuses, not close codes. Granian's RSGI layer uses the value passed to ``close()`` before ``accept()`` as the HTTP status of the refused upgrade; a WebSocket close code like 4401 is not a valid status and panics the worker (``InvalidStatusCode``, and with ``panic = "abort"`` the whole app dies). Drive the rejections through the app with a recording socket so the invariant is checked without a real server. """ @async_test async def test_pre_accept_rejections_are_http_statuses(self) -> None: app, platform, _ = make_platform() session = GameSession( id="seated-game", game_type="dummy", join_code="SEAT01", creator_sub="alice", players=[Seat(user_sub="alice", display_name="Alice")], ) engine = platform.registry.require("dummy") engine.create(session, {}) await platform.game_store.save(session) cases = [ ("unauthenticated", [], "/ws/games/whatever", 401), ("unknown game", [make_user("alice")], "/ws/games/no-such-game", 404), ("not seated", [make_user("mallory")], f"/ws/games/{session.id}", 403), ] for label, users, path, expected in cases: ws = RecordingSocket(path) with ws_users(users): await app.handle_websocket(ws) self.assertEqual([expected], ws.close_codes, label) for code in ws.close_codes: self.assertTrue(100 <= code <= 599, f"{label}: {code} is not an HTTP status") if __name__ == "__main__": unittest.main()