tiangolo.fastapi/tests/test_sse_response.py

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

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


def test_format_sse_with_string_data():
    """Test format_sse with simple string data."""
    result = format_sse("Hello")
    assert result == "data: Hello\n\n"


def test_format_sse_with_dict_data():
    """Test format_sse with dict data."""
    result = format_sse({"data": "Hello", "event": "message"})
    assert result == "event: message\ndata: Hello\n\n"


def test_format_sse_with_all_params():
    """Test format_sse with all parameters."""
    result = format_sse("Hello", event="greeting", id="1", retry=5000)
    assert result == "id: 1\nevent: greeting\nretry: 5000\ndata: Hello\n\n"


def test_format_sse_multiline_data():
    """Test format_sse with multiline data."""
    result = format_sse("Line1\nLine2\nLine3")
    assert result == "data: Line1\ndata: Line2\ndata: Line3\n\n"


def test_format_sse_empty_data():
    """Test format_sse with empty data."""
    result = format_sse("")
    assert result == "data: \n\n"


def test_format_sse_with_event_only():
    """Test format_sse with only event type."""
    result = format_sse(data="test", event="myevent")
    assert result == "event: myevent\ndata: test\n\n"


def test_format_sse_with_retry():
    """Test format_sse with retry."""
    result = format_sse(data="test", retry=3000)
    assert result == "retry: 3000\ndata: test\n\n"


def test_format_sse_with_id():
    """Test format_sse with id."""
    result = format_sse(data="test", id=42)
    assert result == "id: 42\ndata: test\n\n"


# Integration tests with FastAPI


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

    @app.get("/events")
    async def get_events():
        async def event_generator():
            for i in range(3):
                yield {"data": f"Message {i}"}

        return SSEResponse(event_generator())

    client = TestClient(app)
    response = client.get("/events")
    assert response.status_code == 200
    assert response.headers["content-type"] == "text/event-stream; charset=utf-8"
    assert "data: Message 0\n\n" in response.text
    assert "data: Message 1\n\n" in response.text
    assert "data: Message 2\n\n" in response.text


def test_sse_response_with_event_type():
    """Test SSE response with custom event type."""
    app = FastAPI()

    @app.get("/events")
    async def get_events():
        async def event_generator():
            for i in range(2):
                yield {"event": "notification", "data": f"Notice {i}"}

        return SSEResponse(event_generator())

    client = TestClient(app)
    response = client.get("/events")
    assert response.status_code == 200
    assert "event: notification\ndata: Notice 0\n\n" in response.text
    assert "event: notification\ndata: Notice 1\n\n" in response.text


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

    @app.get("/events")
    async def get_events():
        async def event_generator():
            for i in range(2):
                yield {"data": f"Message {i}"}

        return SSEResponse(event_generator(), retry=5000)

    client = TestClient(app)
    response = client.get("/events")
    assert response.status_code == 200
    # Retry is sent as first event
    assert "retry: 5000\n\n" in response.text


def test_sse_response_with_event_serializer():
    """Test SSE response with custom event serializer."""
    app = FastAPI()

    @app.get("/events")
    async def get_events():
        async def simple_generator():
            for i in range(3):
                yield f"Message {i}"

        def serializer(data: str) -> dict[str, str]:
            return {"data": data}

        return SSEResponse(simple_generator(), event_serializer=serializer)

    client = TestClient(app)
    response = client.get("/events")
    assert response.status_code == 200
    assert "data: Message 0\n\n" in response.text
    assert "data: Message 1\n\n" in response.text
    assert "data: Message 2\n\n" in response.text


def test_sse_response_with_sync_generator():
    """Test SSE response with synchronous generator."""
    app = FastAPI()

    @app.get("/events")
    def get_events():
        def sync_generator():
            for i in range(3):
                yield {"data": f"Sync {i}"}

        return SSEResponse(sync_generator())

    client = TestClient(app)
    response = client.get("/events")
    assert response.status_code == 200
    assert "data: Sync 0\n\n" in response.text
    assert "data: Sync 1\n\n" in response.text
    assert "data: Sync 2\n\n" in response.text


def test_sse_response_custom_status_code():
    """Test SSE response with custom status code."""
    app = FastAPI()

    @app.get("/events")
    async def get_events():
        async def event_generator():
            yield {"data": "test"}

        return SSEResponse(event_generator(), status_code=201)

    client = TestClient(app)
    response = client.get("/events")
    assert response.status_code == 201


def test_sse_response_custom_headers():
    """Test SSE response with custom headers."""
    app = FastAPI()

    @app.get("/events")
    async def get_events():
        async def event_generator():
            yield {"data": "test"}

        return SSEResponse(
            event_generator(), headers={"X-Custom-Header": "custom-value"}
        )

    client = TestClient(app)
    response = client.get("/events")
    assert response.status_code == 200
    assert response.headers["x-custom-header"] == "custom-value"


def test_sse_export():
    """Test that SSEResponse is exported from fastapi."""
    from fastapi import SSEResponse, format_sse

    assert SSEResponse is not None
    assert format_sse is not None