Testing interaction of WebSockets with TestClient imposes thread safety?

[jannschu]

Hi,

I'd like to test the interaction between multiple WebSockes using the TestClient and I run into issues that seem to be related with two event loops being created (or threads). On a first glance it looks that my code would need to be thread safe to be usable with two WebSocket test clients in parallel, a requirement that in my case only would be needed for the tests then.

Below is a small example that shows what I'd like to do. A websocket is created and if a second one is created the first shall be killed.

If I run this using Python 3.11 on macOS the program hangs. The expected output would be CancellationError from the first WebSocket. It works if I run it single-threaded using uvicorn. By removing various ingredients like the TaskGroup or the asyncio.Event it also works again hinting at a thread-safety issue.

Should TestClient be implemented to use one thread/run loop only?

Example

import asyncio
import uvicorn

from starlette.testclient import TestClient
from starlette.websockets import WebSocket


lock = asyncio.Lock()  # not thread-safe
current = None


async def run(websocket):
    # removing this stops the hanging
    async with asyncio.TaskGroup() as tg:
        tg.create_task(websocket.receive())


async def app(scope, receive, send) -> None:
    if scope["type"] != "websocket":
        return

    global current

    done = asyncio.Event()  # not thread-safe

    websocket = WebSocket(scope, receive=receive, send=send)
    await websocket.accept()
    task = asyncio.create_task(run(websocket))
    async with lock:
        if current:
            current[0].cancel()
            await current[1].wait()
        current = (task, done)
    try:
        await task
        await websocket.close()
    finally:
        done.set()


# Works with uvicorn
# uvicorn.run(app, port=5001)

client = TestClient(app, backend="asyncio")
with client.websocket_connect("") as ws1, client.websocket_connect("") as ws2:
    pass

[jannschu]

After reading the source code, this seems to work:

client = TestClient(app, backend="asyncio")
with anyio.from_thread.start_blocking_portal(**client.async_backend) as portal:
    client.portal = portal
    with client.websocket_connect("") as ws1, client.websocket_connect("") as ws2:
        pass

Feels a bit hacky. Maybe _portal_factory could reuse a previously created portal?

[Kludex]

This doesn't hang anymore on the latest Starlette.