69 lines
2.4 KiB
Python
69 lines
2.4 KiB
Python
"""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
|