Testing server sent events that use the channels plugin

[aranvir]

Hi, I am a bit stuck and maybe I'm just missing the obvious.

For this I used:

• Python 3.11.7
• litestar 2.8.3
• httpx 0.27.0
• httpx_sse 0.4.0

I have an SSE endpoint which shall forward data from a ChannelPlugin topic. In the example below, I use another endpoint to publish on this topic but this is mostly for demonstrations sake when testing it with an httpx script.

I also included the unit test from https://github.com/litestar-org/litestar/blob/main/tests/unit/test_response/test_sse.py#L16 just to check that I'm not crazy. Anyway.

The SSE endpoint implementation follows https://github.com/orgs/litestar-org/discussions/3104#discussioncomment-8453195

When running the script for testing it works as expected:

• I can connect to the SSE endpoint
• Each message published with the POST request is received by the SSE endpoint

However - I did observe that there is a delay on the SSE endpoint. It takes about 3 seconds for the subscriber to connect to the SSE endpoint and while this happens other requests are also blocked, e.g., the POST request only returns AFTER the SSE connection has been established.

2024-06-05 15:50:06.488 | INFO     | __main__:subscribe:12 - Subscribing to example
2024-06-05 15:50:06.489 | INFO     | __main__:<module>:28 - Started subscription thread.
2024-06-05 15:50:06.489 | INFO     | __main__:<module>:30 - Start sending messages.
2024-06-05 15:50:07.489 | INFO     | __main__:<module>:34 - Sending message
2024-06-05 15:50:09.537 | INFO     | __main__:subscribe:19 - Connected to event source.
2024-06-05 15:50:09.550 | INFO     | __main__:<module>:36 - 201
2024-06-05 15:50:09.551 | INFO     | __main__:subscribe:21 - ServerSentEvent(event='demo', data="{'hello': 'world'}", retry=1000)
...

But at least it works... when I run pytest, it just gets stuck and never finishes. I used the debugger to figure out where it gets stuck and its deep in the httpx code where it waits on transport.handle_request(request) and never returns. So, it might be a question for the httpx repo. I also tested this with an async pytest set up. I also tested it by using threading or asyncio tasks to somehow circumvent the "getting stuck" thing - but neither works

I was also wondering if I'm just doing something fundamentally wrong with setting up the SSE endpoint. I just find it weird that it DOES work without pytest, even though it has this delay. I am aware that subscriber.iter_events() just blocks when there is nothing to retrieve but since it DOES work with the standalone script I had expected that it also somehow would work with pytest.

App

from typing import Iterator

from litestar import Litestar, get, post
from litestar.channels import ChannelsPlugin
from litestar.response import ServerSentEvent
from litestar.channels.backends.memory import MemoryChannelsBackend


@get("/favicon.ico", exclude_from_auth=True)
async def favicon() -> None:
    """Dummy handler to supress error messages in debug mode"""
    ...


@get("/example-sse")
async def example_sse() -> ServerSentEvent:
    def numbers(minimum: int, maximum: int) -> Iterator[str]:
        for i in range(minimum, maximum + 1):
            yield str(i)

    generator = numbers(1, 5)

    return ServerSentEvent(content=generator, event_id="123", event_type="special", retry_duration=1000)


@get("/{topic:str}/notification")
async def notification(topic: str, channels: ChannelsPlugin) -> ServerSentEvent:
    async def generator():
        async with channels.start_subscription([topic]) as subscriber:
            async for event in subscriber.iter_events():
                yield event

    return ServerSentEvent(generator(), event_type="demo", retry_duration=1000)


@post("/{topic:str}/notify")
async def notify(topic: str, data: dict, channels: ChannelsPlugin) -> None:
    channels.publish(str(data), [topic])


channels_plugin = ChannelsPlugin(backend=MemoryChannelsBackend(), arbitrary_channels_allowed=True)


app = Litestar(
    route_handlers=[
        favicon,
        example_sse,
        notification,
        notify
    ],
    plugins=[
        channels_plugin,
    ]
)

Pytest

import threading
from typing import Iterator
import time

import pytest
from litestar.channels import ChannelsPlugin
from litestar.testing import TestClient
from httpx_sse import EventSource, connect_sse

from app import app


@pytest.fixture(name="client")
def fx_client() -> Iterator[TestClient]:
    with TestClient(app=app) as client:
        yield client


@pytest.fixture(name="channels")
def fx_channels(client: TestClient) -> ChannelsPlugin:
    return client.app.plugins.get(ChannelsPlugin)


def test_example_sse(client: TestClient):
    with connect_sse(client, "GET", f"{client.base_url}/example-sse") as event_source:
        events = [sse for sse in event_source.iter_sse()]
        assert len(events) == 5
        for idx, sse in enumerate(events, start=1):
            assert sse.event == "special"
            assert sse.data == str(idx)
            assert sse.id == "123"
            assert sse.retry == 1000


def test_notification(client: TestClient, channels: ChannelsPlugin):
    topic = "demo"
    message = {"hello": "world"}
    with connect_sse(
        client,
        "GET",
        f"{client.base_url}/{topic}/notification",
        timeout=1
    ) as event_source:
        channels.publish(message, [topic])
        received = next(event_source.iter_sse())
        print(received)

Script for testing

import threading
import time

import httpx
import httpx_sse
from loguru import logger

url = "http://localhost:8001"


def subscribe(topic: str):
    logger.info(f"Subscribing to {topic}")
    with httpx.Client() as client:
        with httpx_sse.connect_sse(
            client,
            "GET",
            f"{url}/{topic}/notification"
        ) as event_source:
            logger.info("Connected to event source.")
            for event in event_source.iter_sse():
                logger.info(event)


if __name__ == "__main__":
    topic = "example"
    subscription = threading.Thread(target=subscribe, args=(topic,), daemon=True)
    subscription.start()
    logger.info("Started subscription thread.")
    time.sleep(5)
    logger.info("Start sending messages.")
    with httpx.Client() as client:
        message = {"hello": "world"}
        for _ in range(5):
            logger.info("Sending message")
            response = client.post(f"{url}/{topic}/notify", json=message)
            logger.info(response.status_code)
            time.sleep(0.5)

[euri10]

I notice your using a sync client on an async endpoint, could you test with an async client ?

[aranvir]

As mentioned, same behaviour for sync test as for async test, but find the test code below

import asyncio
import sys
from collections.abc import AsyncIterator

import pytest
from litestar.channels import ChannelsPlugin
from litestar.testing import AsyncTestClient, TestClient
from httpx_sse import aconnect_sse, EventSource, connect_sse

from app import app

if sys.platform == "win32":
    asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())

pytestmark = pytest.mark.anyio


@pytest.fixture
def anyio_backend() -> str:
    return "asyncio"


@pytest.fixture(name="client")
async def fx_client() -> AsyncIterator[AsyncTestClient]:
    async with AsyncTestClient(app=app) as client:
        yield client


@pytest.fixture(name="channels")
def fx_channels(client: AsyncTestClient) -> ChannelsPlugin:
    return client.app.plugins.get(ChannelsPlugin)


async def test_example_sse(client: AsyncTestClient):
    async with aconnect_sse(client, "GET", f"{client.base_url}/example-sse") as event_source:
        events = [sse async for sse in event_source.aiter_sse()]
        assert len(events) == 5
        for idx, sse in enumerate(events, start=1):
            assert sse.event == "special"
            assert sse.data == str(idx)
            assert sse.id == "123"
            assert sse.retry == 1000


async def test_notification(client: AsyncTestClient, channels: ChannelsPlugin):
    topic = "demo"
    message = {"hello": "world"}
    async with aconnect_sse(
        client,
        "GET",
        f"{client.base_url}/{topic}/notification",
        timeout=1
    ) as event_source:
        channels.publish(message, [topic])
        received = anext(event_source.aiter_sse())
        print(received)

[euri10]

spent a little time on this, and yeah it's blocking somewhere in the channels publishing stack, couldnt understand why though

[aranvir]

Thanks for checking ;)

[provinzkraut]

Here's a minimal reproducer. Doesn't seem to be related to channels in particular, but how our tests client handles streams:

import asyncio

from litestar import get, Litestar
from litestar.response import Stream
from litestar.testing import create_async_test_client

async def main():
    @get("/")
    async def stream() -> Stream:
        async def generator():
            while True:
                yield "foo"
                await asyncio.sleep(0.1)

        return Stream(generator())

    async with create_async_test_client([stream]) as client, client.stream("GET", "/") as s:
        await anext(s.aiter_raw())


asyncio.run(main())

When you wrap the same example up in a regular app and run a normal HTTPX client against it, it works as expected. The reason this doesn't work in our test setup is that our test client seems to attempt to consume the whole stream before finishing off the response, which of course doesn't work here.

[aranvir]

Thanks for the clarification! So for now just something I need to test manually or figure out something that does not use the litestar test client.

[provinzkraut]

I'm afraid you'll have to come up with a different way of testing this, because HTTPX (and probably all clients that support calling directly into ASGI apps) will have the same problem.

I'll try to illustrate with a simple example that has nothing to do with Litestar:

import asyncio
import httpx

EVENT = asyncio.Event()


async def app(scope, receive, send):
    await send(
        {
            "type": "http.response.start",
            "status": 200,
            "headers": [
                [b"content-type", b"text/plain"],
            ],
        }
    )
    while not EVENT.is_set():
        await send(
            {
                "type": "http.response.body",
                "body": b"foo",
                "more_body": True,
            }
        )
        await asyncio.sleep(0.1)


async def main():
    async with httpx.AsyncClient(app=app) as client:
        # at this point, HTTPX calls into the ASGI app using 'await app()'
        async with client.stream("GET", "http://localhost") as stream:
            response = await anext(stream.iter_raw())
            EVENT.set()


asyncio.run(main())

Since in the end, ASGI apps are just a single async callable, this callable will have to be awaited somewhere. But as soon as we call await app(...), we are in an infinite loop that we can't move on from without outside interference. Even though we are trying to stream the response body, the app will never stop generating that response body, preventing us from ever reaching the code that would make it stop doing so.

But why does it work with WebSockets?

A similar pattern works with WebSockets and the built-in test client:

import asyncio

from litestar import websocket, WebSocket
from litestar.testing import create_async_test_client

EVENT = asyncio.Event()


@websocket("/")
async def handler(socket: WebSocket) -> None:
    await socket.accept()
    while not EVENT.is_set():
        await socket.send_text("foo")
        await asyncio.sleep(0.1)
    await socket.close()


async def main():
    async with create_async_test_client([handler]) as client:
        with await client.websocket_connect("/") as ws:
            assert ws.receive_text() == "foo"
            EVENT.set()


asyncio.run(main())

The reason this works is because, for WebSockets, Litestar is cheating. It's not doing a full WebSocket roundtrip, and neither is the underlying HTTPX test client. We simply create an internal queue where the messages are stored, from which the client can then read them.

---

I'm not sure what the best way to handle this would be though. Maybe we create another test client that has a fully fledged uvicorn running behind the scenes in a subprocess and uses a regular HTTPX client to connect to it? The other option would be faking the HTTP layer for these requests, but I don't really like this idea.

@litestar-org/members, thoughts?

[aranvir]

Thanks for the insights!!

I'm gonna look into uvicorn subprocess + httpx for testing my SSE endpoints.

[provinzkraut]

We have something like this in place already, but only use it to build a part of our docs:

litestar/tools/sphinx_ext/run_examples.py

Line 62 in 3a6a293

def run_app(path: Path) -> Generator[int, None, None]:

If you're interested in contributing, a PR that adds a test client based on this pattern would definitely be welcome :)

[aranvir]

Added a PR #3655

[aranvir]

Tinkered a bit with the example provided by @provinzkraut and came up with the workaround shared below. I wanted to share this here first for anyone interested before starting on a PR. Also to get some early feedback if the helper module fits into the litestar testing library.

For my use case the solution is good enough, but it breaks with the approach to provide test clients. It also splits the test into an app module where I set up the components that I want to tests and into a test module where I actually run the tests. Maybe there is a way to combine this for a smooth user experience - otherwise, it could be used as a pattern for "true" deployed testing instead of running the ASGI app directly in the httpx client.

Let me know what you think!

Description

I could not get the multiprocessing approach to work with my application - I always got a "pickling" error when it came to the app object and I could not resolve it. Since the object is loaded from a file via importlib anyway I thought - why not just run the server as a subprocess from that file. So that's what helper.run_app does: Start a webserver via litestar CLI in a subprocess.

The app is a minimal example with SSE. I'm using Redis here as the channels backend, so I can publish messages from the test client. Ideally, both app and test client get the ChannelsPlugin from the same source - for demonstration purposes I just copy-pasted it. If I had gotten a solution with multiprocessing to work, adding another backend with a pipe might have been attractive to circumvent the need for Redis or another broker - but I find this also acceptable.

Lastly, the test module: Two fixtures - one to get the server running and return the url, the other one to set get a ChannelsPlugin instance so we can generate messages. Within the test, I use an async task to periodically publish messages. Then we connect to the SSE endpoint and after the first message has been received and verified the test exits.

Helper module

"""
Provide a utility function to run an app as a subprocess.
"""
import pathlib
import subprocess
import time
import socket
from contextlib import contextmanager

import httpx


class StartupError(RuntimeError):
    pass


def _get_available_port() -> int:
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
        # Bind to a free port provided by the host
        try:
            sock.bind(("localhost", 0))
        except OSError as e:
            raise StartupError("Could not find an open port") from e
        else:
            return sock.getsockname()[1]


@contextmanager
def run_app(workdir: pathlib.Path, app: str) -> str:
    port = _get_available_port()
    proc = subprocess.Popen(
        args=[
            "litestar",
            "--app",
            app,
            "run",
            "--port",
            str(port)
        ],
        stderr=subprocess.PIPE,
        stdout=subprocess.PIPE,
        cwd=workdir
    )
    url = f"http://127.0.0.1:{port}"
    for _ in range(100):
        try:
            httpx.get(url, timeout=0.1)
            break
        except httpx.TransportError:
            time.sleep(1)
    yield url
    proc.kill()

App module

"""
Assemble components into an app that shall be tested
"""
from litestar import Litestar, get
from litestar.channels import ChannelsPlugin
from litestar.channels.backends.redis import RedisChannelsPubSubBackend
from litestar.response import ServerSentEvent
from redis.asyncio import Redis


@get("/notify/{topic:str}")
async def get_notified(
        topic: str,
        channels: ChannelsPlugin
) -> ServerSentEvent:
    async def generator():
        async with channels.start_subscription([topic]) as subscriber:
            async for event in subscriber.iter_events():
                yield event

    return ServerSentEvent(generator(), event_type="CommandHistoryNotification")


def create_test_app():
    redis_instance = Redis()
    channels_backend = RedisChannelsPubSubBackend(redis=redis_instance)
    channels_instance = ChannelsPlugin(backend=channels_backend, arbitrary_channels_allowed=True)

    return Litestar(
        route_handlers=[get_notified],
        plugins=[channels_instance]
    )


app = create_test_app()

Test module

"""
Test the app running in a subprocess
"""
import asyncio
import pathlib
import sys

import pytest
import httpx
import httpx_sse

from litestar.channels import ChannelsPlugin
from litestar.channels.backends.redis import RedisChannelsPubSubBackend
from redis.asyncio import Redis

import helper

if sys.platform == "win32":
    asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())

pytestmark = pytest.mark.anyio


@pytest.fixture(scope="session")
def anyio_backend() -> str:
    return "asyncio"


ROOT = pathlib.Path(__file__).parent


@pytest.fixture(name="notification_url", scope="session")
async def fx_notification_url() -> str:
    with helper.run_app(workdir=ROOT, app="example_app:app") as url:
        yield url


@pytest.fixture(name="redis_channels", scope="session")
async def fx_redis_channels():
    redis_instance = Redis()
    channels_backend = RedisChannelsPubSubBackend(redis=redis_instance)
    channels_instance = ChannelsPlugin(backend=channels_backend, arbitrary_channels_allowed=True)
    await channels_instance._on_startup()
    yield channels_instance
    await channels_instance._on_shutdown()


async def test_connection(notification_url: str, redis_channels: ChannelsPlugin):
    topic = "demo"
    message = "hello"

    running = asyncio.Event()
    running.set()

    async def send_notifications():
        print("started task")
        while running.is_set():
            await redis_channels.wait_published(message, channels=[topic])
            await asyncio.sleep(0.5)

    task = asyncio.create_task(send_notifications())

    async with httpx.AsyncClient(base_url=notification_url) as client:
        async with httpx_sse.aconnect_sse(
            client,
            "GET",
            f"{notification_url}/notify/{topic}"
        ) as event_source:
            async for event in event_source.aiter_sse():
                assert event.data == message
                running.clear()
                break

[euri10]

you could also hack this by setting a timeout that default to None in your generator and feed it with some value in your tests, I'm doing that in https://github.com/euri10/litestar_saq_htmx/blob/main/tests/test_saq_htmx.py#L12