Browse Source
Add native SSE support to FastAPI with: - SSEResponse: A StreamingResponse subclass that automatically formats events according to the SSE specification with support for event types, data, IDs, and retry configuration - sse_content: A helper function to create SSE-formatted content from a dictionary iterator Both exports are added to the public FastAPI API. Co-Authored-By: Claude Opus 4.6 <[email protected]>pull/14976/head
3 changed files with 362 additions and 1 deletions
@ -0,0 +1,224 @@ |
|||
"""Tests for Server-Sent Events (SSE) response.""" |
|||
from collections.abc import AsyncGenerator |
|||
from typing import Any |
|||
|
|||
import pytest |
|||
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 |
|||
Loading…
Reference in new issue