Web frontend

python-backend/tests/websocket_harness.py

@@ -0,0 +1,68 @@
+"""Drive an ASGI websocket endpoint on the caller's own event loop.
+
+Starlette's ``TestClient`` runs the app in a second thread with its own loop, which
+would put the bus and the websockets on different loops - so ``bus.emit`` would take
+its thread-safe path and a following ``drain()`` could return before the event was even
+queued. In the real app there is only ever one loop, and this harness keeps the tests
+that way.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import json
+from collections.abc import AsyncIterator
+from typing import Any
+
+
+class WebSocketSession:
+    def __init__(self) -> None:
+        self.to_app: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
+        self.from_app: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
+
+    async def receive(self) -> dict[str, Any]:
+        return await self.to_app.get()
+
+    async def send(self, message: dict[str, Any]) -> None:
+        await self.from_app.put(message)
+
+    async def next_json(self, timeout: float = 2.0) -> dict[str, Any]:
+        """The next ``websocket.send`` frame, decoded."""
+        while True:
+            message = await asyncio.wait_for(self.from_app.get(), timeout)
+            if message["type"] == "websocket.send":
+                text: str = message["text"]
+                return json.loads(text)
+
+
+@contextlib.asynccontextmanager
+async def websocket_connect(app: Any, path: str) -> AsyncIterator[WebSocketSession]:
+    session = WebSocketSession()
+    scope = {
+        "type": "websocket",
+        "asgi": {"version": "3.0", "spec_version": "2.3"},
+        "http_version": "1.1",
+        "scheme": "ws",
+        "path": path,
+        "raw_path": path.encode(),
+        "query_string": b"",
+        "root_path": "",
+        "headers": [(b"host", b"testserver")],
+        "client": ("testclient", 50000),
+        "server": ("testserver", 80),
+        "subprotocols": [],
+        "state": {},
+    }
+    await session.to_app.put({"type": "websocket.connect"})
+    task = asyncio.create_task(app(scope, session.receive, session.send))
+
+    accepted = await asyncio.wait_for(session.from_app.get(), 2.0)
+    assert accepted["type"] == "websocket.accept", accepted
+    try:
+        yield session
+    finally:
+        await session.to_app.put({"type": "websocket.disconnect", "code": 1000})
+        task.cancel()
+        with contextlib.suppress(asyncio.CancelledError, Exception):
+            await task