WebSocket server with message handler and worker

client.py

import time

import trio
from trio_websocket import open_websocket_url


async def send_one_message(i):
    try:
        async with open_websocket_url("ws://127.0.0.1:8000/foo") as ws:
            await ws.send_message(f"hello, {i}!")
            start = time.perf_counter()
            message = await ws.get_message()
            end = time.perf_counter()
            print(f"Received message: '{message}' in {end-start:.0f}s")
    except OSError as ose:
        print(f"Connection attempt failed: {ose}")


async def main():
    async with trio.open_nursery() as nursery:
        for i in range(10):
            nursery.start_soon(send_one_message, i)


trio.run(main)

server.py

from functools import partial
import trio
from trio_websocket import serve_websocket, ConnectionClosed


async def worker(ws, message):
    await trio.sleep(2)
    await ws.send_message(f"work on {message} done")


async def message_handler(request, *, nursery):
    ws = await request.accept()
    while True:
        try:
            message = await ws.get_message()
            nursery.start_soon(worker, ws, message)
        except ConnectionClosed:
            break


async def main():
    async with trio.open_nursery() as nursery:
        await serve_websocket(
            partial(message_handler, nursery=nursery),
            "127.0.0.1",
            8000,
            ssl_context=None,
        )


trio.run(main)