async def test_cannot_nest_contexts(): dispatch = Dispatch() class MyRequestContext: x: int = 0 context = MyRequestContext() with pytest.raises(RuntimeError): async with dispatch.connection_context(context): async with dispatch.connection_context(context): pass
async def test_handler_raises_jsonrpc_exc(): """ A JSON-RPC exception in a handler should be raised. """ dispatch = Dispatch() @dispatch.handler async def foo_bar(*args, **kwargs): raise JsonRpcInternalError("foo bar") with pytest.raises(JsonRpcInternalError): result = await dispatch.execute(JsonRpcRequest(id=0, method="foo_bar"))
async def test_handler_raises_exc(caplog): """ A non-JSON-RPC exception in a handler should be raised and logged. """ dispatch = Dispatch() @dispatch.handler async def foo_bar(*args, **kwargs): raise Exception("foo bar") with pytest.raises(JsonRpcInternalError): await dispatch.execute(JsonRpcRequest(id=0, method="foo_bar")) assert 'An unhandled exception occurred in handler "foo_bar"' in caplog.text
async def test_context_fails_if_not_set(caplog): """ If a handler accesses the connection context but the context was never set, then it should raise a RuntimeError. Since it occurs in a handler, execute() converts it to an internal error and also logs it. """ dispatch = Dispatch() @dispatch.handler async def increment(): dispatch.ctx.x += 1 return {"x": dispatch.ctx.x} with pytest.raises(JsonRpcInternalError): await dispatch.execute(JsonRpcRequest(id=0, method="increment")) assert ("The .context property is only valid in a connection context" in caplog.text)
async def test_dispatch_context_inside_request(): """ If a context is set, it should be visible inside the handler and the handler should be able to modify the state in ways that are visible to other handlers. """ dispatch = Dispatch() class MyRequestContext: x: int = 0 @dispatch.handler async def increment(): dispatch.ctx.x += 1 return {"x": dispatch.ctx.x} context = MyRequestContext() async with dispatch.connection_context(context): result = await dispatch.execute( JsonRpcRequest(id=0, method="increment")) assert result["x"] == 1 result = await dispatch.execute( JsonRpcRequest(id=1, method="increment")) assert result["x"] == 2 assert context.x == 2
async def test_dispatch_requests(): dispatch = Dispatch() h1_calls = 0 h1_args = None h2_calls = 0 h2_args = None @dispatch.handler async def handler1(*args, **kwargs): nonlocal h1_calls nonlocal h1_args h1_calls += 1 h1_args = args, kwargs return {"handled_by": "handler1"} @dispatch.handler async def handler2(*args, **kwargs): nonlocal h2_calls nonlocal h2_args h2_calls += 1 h2_args = args, kwargs return {"handled_by": "handler2"} req1 = JsonRpcRequest(id=0, method="handler1", params={"foo": "bar"}) result1 = await dispatch.execute(req1) assert result1["handled_by"] == "handler1" assert h1_calls == 1 assert h2_calls == 0 assert h1_args == (tuple(), {"foo": "bar"}) req2 = JsonRpcRequest(id=0, method="handler2", params=[1, 2, 3]) result2 = await dispatch.execute(req2) assert result2["handled_by"] == "handler2" assert h1_calls == 1 assert h2_calls == 1 assert h2_args == ((1, 2, 3), {}) req3 = JsonRpcRequest(id=0, method="handler1") result3 = await dispatch.execute(req3) assert result3["handled_by"] == "handler1" assert h1_calls == 2 assert h2_calls == 1 assert h1_args == (tuple(), {})
JsonRpcException, ) from trio_jsonrpc.transport.ws import WebSocketTransport import trio_websocket from .shared import AuthorizationError, InsufficientFundsError user_pins = { "john": 1234, "jane": 5678, } user_balances = { "john": 100, "jane": 100, } dispatch = Dispatch() logger = logging.getLogger("server") @dataclass class ConnectionContext: """ This object stores the context data for each connection. """ user: typing.Optional[str] = None @dispatch.handler async def login(user: str, pin: int) -> bool: """ Verify the user's pin and update connection context.
async def test_websocket_roundtrip(nursery): """ Create a WebSocket server and several clients. """ dispatch = Dispatch() client_count = 0 @dispatch.handler async def hello_world(n): return {"mode": "hello", "client_number": n} @dispatch.handler async def goodbye_world(n): return {"mode": "goodbye", "client_number": n} async def client(n): nonlocal client_count url = f"ws://localhost:{server_port}" logging.info("Client #%d: Connecting to %s", n, url) async with open_jsonrpc_ws(url) as client_conn: logging.info("Client #%d: Sending hello", n) resp = await client_conn.request(method="hello_world", params={"n": n}) logging.info("Client #%d: Got hello response", n) assert resp["mode"] == "hello" assert resp["client_number"] == n logging.info("Client #%d: Sending goodbye", n) resp = await client_conn.request(method="goodbye_world", params={"n": n}) logging.info("Client #%d: Got goodbye response", n) assert resp["mode"] == "goodbye" assert resp["client_number"] == n client_count += 1 async def responder(recv_channel, conn): async for request, result in recv_channel: if isinstance(result, JsonRpcException): await conn.respond_with_error(request, result) else: await conn.respond_with_result(request, result) async def connection_handler(ws_request): ws = await ws_request.accept() transport = WebSocketTransport(ws) rpc_conn = JsonRpcConnection(transport, JsonRpcConnectionType.SERVER) result_send, result_recv = trio.open_memory_channel(10) async with trio.open_nursery() as conn_nursery: conn_nursery.start_soon(rpc_conn._background_task) conn_nursery.start_soon(responder, result_recv, rpc_conn) logging.info("Serving requests on new connection...") async for request in rpc_conn.iter_requests(): print(request) conn_nursery.start_soon(dispatch.handle_request, request, result_send) conn_nursery.cancel_scope.cancel() server = await nursery.start(trio_websocket.serve_websocket, connection_handler, "localhost", 0, None) server_port = server.port logging.info("Server is listening on port %d", server_port) async with trio.open_nursery() as client_nursery: client_nursery.start_soon(client, 1) client_nursery.start_soon(client, 2) client_nursery.start_soon(client, 3) client_nursery.start_soon(client, 4) client_nursery.start_soon(client, 5) assert client_count == 5
async def test_dispatch_method_not_found(): dispatch = Dispatch() with pytest.raises(JsonRpcMethodNotFoundError): await dispatch.execute(JsonRpcRequest(id=0, method="hello_world"))
async def test_dispatch_requires_named_method(): dispatch = Dispatch() with pytest.raises(RuntimeError): await dispatch.handler(object())