Using httpx-ws as background task to test websockets

[wholmen]

Hi, reposting the issue to this github.

I am struggling with httpx-ws. It works when the server-side websocket connection is closed, but my implementations are based on websockets that are open until the client breaks them. Could you help me adjusting the code to make the implementation work?

Websocket implementation

from fastapi import APIRouter, Depends, WebSocket, WebSocketDisconnect
import asyncio

@router.websocket("/ws")
async def websocket_labtest(websocket: WebSocket):
    await websocket.accept()
    sub_id: uuid.UUID = uuid.uuid4()
    subjects.ADD_LABTEST_QUEUE[sub_id] = asyncio.Queue()

    try:
        while True:
            event: events.LabtestAddedEvent = await subjects.ADD_LABTEST_QUEUE[sub_id].get()
            await websocket.send_json(event.json())
    except WebSocketDisconnect:
        subjects.ADD_LABTEST_QUEUE.pop(sub_id)

Test implementation

@pytest.fixture
async def labtest_client_async() -> AsyncClient:
    async with AsyncClient(transport=ASGIWebSocketTransport(app), base_url='http://test') as client:
        yield client

async def subscribe_to_labtests_ws_async(labtest_client_async: httpx.AsyncClient, queue: asyncio.Queue):
    async with httpx_ws.aconnect_ws("/labtests/ws", labtest_client_async) as ws:
        message = await ws.receive_json()
        queue.put_nowait(message)

@pytest.mark.asyncio
async def test_add_lab_test(labtest_client_async: httpx.AsyncClient):
    sub_queue: asyncio.Queue = asyncio.Queue()
    task = asyncio.create_task(subscribe_to_labtests_ws_async(labtest_client_async, sub_queue))

    try:
        await asyncio.sleep(SHORT_WAIT)
        await add_labtest(labtest_client_async)
        await asyncio.sleep(SHORT_WAIT)
    finally:
        task.cancel()

    message_labtest1 = await sub_queue.get()
    print(f"{message_labtest1=}")

The websocket works fine when I use e.g. Postman to connect to the websocket. It seems httpx_ws is unable to close the connection even though I say task.cancel(). That makes the test run until I force quit.

It seems like the problem being related to asyncio.create_task. If I don't wait for the "add labtest event" in server, I manage to get the message sent to me. So it seems like the connection struggles to run as a background task. Maybe this is completely out of scope. I am open to ideas on how to do this kind of integration testing.

I know the code is not self-contained. If it's hard to understand, I can make a simpler self-contained example

[frankie567]

Hi @wholmen 👋

Is there a reason why you want to run your client connection into a task? IMO, it would be way easier for debugging/reasoning to just run it in the main context:

@pytest.mark.asyncio
async def test_add_lab_test(labtest_client_async: httpx.AsyncClient):
    sub_queue: asyncio.Queue = asyncio.Queue()
    async with httpx_ws.aconnect_ws("/labtests/ws", labtest_client_async) as ws:
        message = await ws.receive_json()
        await sub_queue.put(message)
        await asyncio.sleep(SHORT_WAIT)

    message_labtest1 = await sub_queue.get()
    print(f"{message_labtest1=}")