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=}")

[wholmen]

Yes, I need it to be a background task, because it's the await add_labtest(labtest_client_async) that tells the websocket client to fire a message to the client.

So the full end-to-end system I want to test here is:

1. (client) Start subscription to a webservice
2. (client) Send rest api call to add a lab test to the database
3. (server) adds labtest to database
4. (server) Server puts an event to a queue that will tell every subscriber that a labtest was added
5. (server) Websocket picks up the event and notifies every webservice subscribers
6. (client) Picks up the message that the labtest was created
7. (client) Query API to retrieve data on the new labtest

I followed this pattern with graphql and subscriptions, but I wanted to get rid of graphql for other reasons. So now I'm trying to replace subscriptions with websocket, but it proves a bit difficult to test it with pytest.

The information flow I have described works with manual testing. I'm just struggling to test it with pytest.

I could of course just skip end-to-end-testing, but I like to ensure these more advanced patterns do not break

[wholmen]

I could write a more complete (self-contained) example if the code example I provided is not obvious

[frankie567]

Ok, I see. I think the problem here is that you're using the same HTTPX client instance. I suspect the first one is blocked by the pending WebSocket Client. Could you try to open another one to make your add_labtest request?

[wholmen]

I tried to use a different httpx AsyncClient now for the post request.

The background task still doesn't manage to pick up the websocket message, and it doesn't manage to shut down after the test has passed

[frankie567]

If you could provide a simplified example, that would be great and help me to understand all the bits and pieces.

[wholmen]

Ok, I have been a little busy, but I have created a full example showing the workflow. I have tested this with Postman, so it should in principle be testable.

Implementation of websocket

import asyncio
import uuid
from typing import Dict

from fastapi import APIRouter, Response, status, WebSocket, WebSocketDisconnect


router = APIRouter(tags=["Test e2e of websocket"])

SOCKET_QUEUE: Dict[uuid.UUID, asyncio.Queue[int]] = {}


@router.post("/activate_websocket/{post_id}")
def activate_websocket(post_id: int):
    for socket in SOCKET_QUEUE.keys():
        SOCKET_QUEUE[socket].put_nowait(post_id)

    return Response(status_code=status.HTTP_200_OK)


@router.websocket_route("/test_ws")
async def ws_test(websocket: WebSocket):
    await websocket.accept()
    socket_id: uuid.UUID = uuid.uuid4()
    SOCKET_QUEUE[socket_id] = asyncio.Queue()

    try:
        while True:
            post_id: int = await SOCKET_QUEUE[socket_id].get()
            await websocket.send_text(f"Post id {post_id} was sent to active_websocket")
    except WebSocketDisconnect:
        SOCKET_QUEUE.pop(socket_id)

Test

import pytest
import asyncio
import httpx
import httpx_ws
from httpx import AsyncClient
from httpx_ws.transport import ASGIWebSocketTransport
from fastapi import status

SHORT_WAIT = 0.2


@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_simple_websocket_async(labtest_client_async: httpx.AsyncClient, queue: asyncio.Queue):
    async with httpx_ws.aconnect_ws("http://test/test_ws", labtest_client_async) as ws:
        message = await ws.receive_text()
        queue.put_nowait(message)


async def post_id_to_send_websocket_message(labtest_client_async: httpx.AsyncClient, post_id: int):
    url = f"/activate_websocket/{post_id}"
    response: httpx.Response = await labtest_client_async.post(url)
    assert response.status_code == status.HTTP_200_OK


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

    try:
        await asyncio.sleep(SHORT_WAIT)
        await post_id_to_send_websocket_message(labtest_client_async, post_id=2)
        await asyncio.sleep(SHORT_WAIT)
        await post_id_to_send_websocket_message(labtest_client_async, post_id=3)
        await asyncio.sleep(SHORT_WAIT)
    finally:
        task.cancel()

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

Testing this from Postman works fine, and every post-request I send will trigger the websocket:

The reason I want this pattern, is that I want to have items posted to the database from rabbitmq, and then notify the client when it happens. I also want to be able to parse items one by one, and tell the client that things are ready. I am making software that has some heavy computational components

edit: added two imports to make fixture work. It's in conftest.py in my setup, so I made a mistake when copying it here

[wholmen]

You added while True: to the websocket test. That's interesting. Didn't think of that, but it makes sense. Works now :)

Thanks a lot!

[wholmen]

The test don't end properly, but I can look into that myself. Since I might not end up using this pattern after all, you don't have to spend much more time on it :)

[frankie567]

Oh yeah, right indeed! Since you want to receive several messages, the while loop is necessary. And yes, the test doesn't end up properly. You would likely need a mechanism to signal the thread to stop, e.g. using threading.Event.

[wholmen]

Is it possible to close the websocket from the client side using httpx_ws?

A brute force ws.close() seems to do nothing. The background task is ended properly, but none of them does the trick.

[frankie567]

Exiting the context manager is enough to close the client connection, it's implicitly closed when leaving the context manager.

The problem that may happen is that you're stuck waiting for receive_text, so you don't have the opportunity to go back to the start of the loop. You can add a timeout to prevent it from blocking for too long: await ws.receive_text(timeout=0.1)