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())