tiangolo.fastapi/tests/test_sse_response.py

"""Tests for Server-Sent Events (SSE) response."""

from collections.abc import AsyncGenerator
from typing import Any

from fastapi import FastAPI
from fastapi.responses import SSEResponse, sse_content
from fastapi.testclient import TestClient


def test_sse_response_basic():
    """Test basic SSE response with event data."""
    app = FastAPI()

    @app.get("/events")
    def events():
        async def generator() -> AsyncGenerator[dict[str, Any], None]:
            yield {"event": "message", "data": "Hello"}
            yield {"data": "World"}

        return SSEResponse(generator())

    client = TestClient(app, raise_server_exceptions=False)
    response = client.get("/events")

    assert response.status_code == 200
    assert "text/event-stream" in response.headers["content-type"]
    # Check that the SSE format is correct
    content = response.text
    assert "event: message" in content
    assert "data: Hello" in content
    assert "data: World" in content


def test_sse_response_with_retry():
    """Test SSE response with retry configuration."""
    app = FastAPI()

    @app.get("/events")
    def events():
        async def generator() -> AsyncGenerator[dict[str, Any], None]:
            yield {"data": "Hello"}

        return SSEResponse(generator(), retry=5000)

    client = TestClient(app, raise_server_exceptions=False)
    response = client.get("/events")

    assert response.status_code == 200
    assert "retry: 5000" in response.text


def test_sse_response_with_event_id():
    """Test SSE response with event ID."""
    app = FastAPI()

    @app.get("/events")
    def events():
        async def generator() -> AsyncGenerator[dict[str, Any], None]:
            yield {"event": "update", "data": "test", "id": "1"}

        return SSEResponse(generator())

    client = TestClient(app, raise_server_exceptions=False)
    response = client.get("/events")

    assert response.status_code == 200
    assert "event: update" in response.text
    assert "data: test" in response.text
    assert "id: 1" in response.text


def test_sse_response_string_content():
    """Test SSE response with pre-formatted string content."""
    app = FastAPI()

    @app.get("/events")
    def events():
        async def generator() -> AsyncGenerator[str, None]:
            yield "data: Hello\n\n"

        return SSEResponse(generator())

    client = TestClient(app, raise_server_exceptions=False)
    response = client.get("/events")

    assert response.status_code == 200
    assert "data: Hello" in response.text


def test_sse_response_json_data():
    """Test SSE response with JSON data."""
    app = FastAPI()

    @app.get("/events")
    def events():
        async def generator() -> AsyncGenerator[dict[str, Any], None]:
            yield {"data": {"message": "hello", "count": 42}}

        return SSEResponse(generator())

    client = TestClient(app, raise_server_exceptions=False)
    response = client.get("/events")

    assert response.status_code == 200
    # JSON data should be serialized
    assert '"message": "hello"' in response.text
    assert '"count": 42' in response.text


def test_sse_response_with_retry_in_event():
    """Test SSE response retry in event data overrides default."""
    app = FastAPI()

    @app.get("/events")
    def events():
        async def generator() -> AsyncGenerator[dict[str, Any], None]:
            # Event-level retry should override constructor retry
            yield {"data": "Hello", "retry": 3000}

        return SSEResponse(generator(), retry=5000)

    client = TestClient(app, raise_server_exceptions=False)
    response = client.get("/events")

    assert response.status_code == 200
    assert "retry: 3000" in response.text
    # The default retry should not appear
    lines = response.text.split("\n")
    retry_count = sum(1 for line in lines if line.startswith("retry:"))
    assert retry_count == 1


def test_sse_content_helper():
    """Test the sse_content helper function."""
    app = FastAPI()

    @app.get("/events")
    def events():
        async def generator() -> AsyncGenerator[dict[str, Any], None]:
            yield {"event": "message", "data": "Hello"}
            yield {"data": "World"}

        from starlette.responses import StreamingResponse

        return StreamingResponse(
            sse_content(generator()), media_type="text/event-stream"
        )

    client = TestClient(app, raise_server_exceptions=False)
    response = client.get("/events")

    assert response.status_code == 200
    assert "text/event-stream" in response.headers["content-type"]
    content = response.text
    assert "event: message" in content
    assert "data: Hello" in content
    assert "data: World" in content


def test_sse_response_custom_media_type():
    """Test SSE response with custom media type."""
    app = FastAPI()

    @app.get("/events")
    def events():
        async def generator() -> AsyncGenerator[dict[str, Any], None]:
            yield {"data": "Hello"}

        return SSEResponse(generator(), media_type="text/event-stream; charset=utf-8")

    client = TestClient(app, raise_server_exceptions=False)
    response = client.get("/events")

    assert response.status_code == 200
    assert "text/event-stream" in response.headers["content-type"]


def test_sse_response_data_newlines():
    """Test SSE response with multiline data."""
    app = FastAPI()

    @app.get("/events")
    def events():
        async def generator() -> AsyncGenerator[dict[str, Any], None]:
            yield {"data": "line1\nline2\nline3"}

        return SSEResponse(generator())

    client = TestClient(app, raise_server_exceptions=False)
    response = client.get("/events")

    assert response.status_code == 200
    # Each line should be prefixed with "data: "
    assert "data: line1" in response.text
    assert "data: line2" in response.text
    assert "data: line3" in response.text


def test_sse_response_multiple_events():
    """Test SSE response with multiple events."""
    app = FastAPI()

    @app.get("/events")
    def events():
        async def generator() -> AsyncGenerator[dict[str, Any], None]:
            yield {"event": "first", "data": "data1"}
            yield {"event": "second", "data": "data2", "id": "2"}
            yield {"data": "data3"}

        return SSEResponse(generator())

    client = TestClient(app, raise_server_exceptions=False)
    response = client.get("/events")

    assert response.status_code == 200
    content = response.text
    # First event
    assert "event: first" in content
    assert "data: data1" in content
    # Second event
    assert "event: second" in content
    assert "data: data2" in content
    assert "id: 2" in content
    # Third event (no event type)
    assert content.count("data: data3") == 1