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