Testing with Channels and Websockets

[TheZwieback]

Hi,
I currently have an issue with testing a channels and websocket setup very similar to the docs. I create a TestClient and open a websocket_connection, but when sending a message via channels the receive_text() either blocks and prevents delivery of the message via channels or is skipped when I set block to False. I tried with a while True: loop but that did not work.
A second issue is that I cannot close the connection so the test runs forever. Is there a way to test this?

from litestar import Litestar, websocket, WebSocket
from litestar.testing import TestClient
from litestar.channels import ChannelsPlugin
from litestar.channels.backends.memory import MemoryChannelsBackend
from litestar.channels import ChannelsPlugin
import logging


@websocket("/ws")
async def ws(socket: WebSocket, channels: ChannelsPlugin) -> None:
    await socket.accept()

    async with channels.start_subscription(["my_new_channel"]) as subscriber:
        async for message in subscriber.iter_events():
            logging.info(message)
            await socket.send_text(message)


channels = ChannelsPlugin(
    backend=MemoryChannelsBackend(history=0), channels=["my_new_channel"], arbitrary_channels_allowed=True
)

with TestClient(
    app=Litestar(
        debug=True,
        plugins=[channels],
        route_handlers=[ws],
    )
) as client:
    with client.websocket_connect("/ws") as websocket_connection:
        logging.info(channels)
        channels.publish("asdf", "my_new_channel")
        message = websocket_connection.receive_text(block=False)
        logging.info(message)
        websocket_connection.close()

[aranvir]

Hey, the code below should work. So what's different:

I tried to figure out where your code was blocking - and it blocked when opening the websocket connection. This happens because it gets stuck waiting for subscriber.iter_events(). As mentioned in the docs using run_in_background is preferred, since this avoids the blocking in the foreground.

However, we need to use "something" blocking to keep the websocket connection open. I decided to use socket.recv_text() as an example because it's an async block so it allows the rest of the application to keep running. However, in this example this would also mean that the connection would close as soon as the socket actually receives a message. So there are also other options how to deal with this.. as long as it's an "async block".

Hope that helps :)

from litestar import Litestar, websocket, WebSocket
from litestar.testing import create_test_client
from litestar.channels.backends.memory import MemoryChannelsBackend
from litestar.channels import ChannelsPlugin
import logging


def test_websocket():
    @websocket("/ws")
    async def ws(socket: WebSocket, channels: ChannelsPlugin) -> None:
        await socket.accept()

        async def _sender(message: str):
            await socket.send_text(message)

        async with channels.start_subscription(["my_new_channel"]) as subscriber:
            async with subscriber.run_in_background(_sender):
                await socket.receive_text()

    channels = ChannelsPlugin(
        backend=MemoryChannelsBackend(history=0), channels=["my_new_channel"], arbitrary_channels_allowed=True
    )

    with create_test_client(
        route_handlers=[ws],
        plugins=[channels]
    ) as client:
        conn = client.websocket_connect("/ws")
        with conn:
            channels.publish("asdf", "my_new_channel")
            message = conn.receive_text(timeout=1)
            logging.info(message)

[TheZwieback]

Thanks for the help. I did something similar in the end with an open async future, still I could get it not tested properly as the sync context blocked the async execution. In the end I gave up trying to get this to work with the sync client (maybe the async would work but the websocket seems to be broken there). Instead I just wrote an integration test, basically just starting an app and doing some black box tests. Not great but at least some testing.

[aranvir]

You might have to check the litestar version you are using, since the async websocket has only been added with 2.9.0 to the test client. I'm pretty sure I used both the sync and the async websocket test client in some of my test code and both worked.

[TheZwieback]

Seems if I use the helper method the AsyncTestClient works with the WebSocket :). But looks like the WebSocket stays synchronous. I was hoping I could put the WebSocket context to sleep after publishing or something like that. I think without some sleep / pause to let the background or WebSocket loop take over there is no good way to test this via the client. API test it is.

from litestar import Litestar, websocket, WebSocket
from litestar.testing import create_async_test_client
from litestar.channels.backends.memory import MemoryChannelsBackend
from litestar.channels import ChannelsPlugin
import logging
import asyncio


async def test_websocket():
    @websocket("/ws")
    async def ws(socket: WebSocket, channels: ChannelsPlugin) -> None:
        await socket.accept()

        async def _sender(message: str):
            await socket.send_text(message)

        async with channels.start_subscription(["my_new_channel"]) as subscriber:
            async for message in subscriber.iter_events():
                logging.info(message)
                await socket.send_text(message)

    channels = ChannelsPlugin(
        backend=MemoryChannelsBackend(history=0), channels=["my_new_channel"], arbitrary_channels_allowed=True
    )

    async with create_async_test_client(route_handlers=[ws], plugins=[channels]) as client:
        conn = await client.websocket_connect("/ws")
        with conn:
            channels.publish("asdf", "my_new_channel")
            await asyncio.sleep(1)
            message = conn.receive_text(timeout=1)
            logging.info(message)


loop = asyncio.get_event_loop()
loop.run_until_complete(test_websocket())
loop.close()