Browse Source
Add SSEResponse class and format_sse function for Server-Sent Events support in FastAPI responses. Supports both async and sync generators, event serialization, retry configuration, and automatic SSE formatting. Co-Authored-By: Claude Opus 4.6 <[email protected]>pull/14977/head
3 changed files with 375 additions and 0 deletions
@ -0,0 +1,201 @@ |
|||||
|
"""Tests for Server-Sent Events (SSE) response classes.""" |
||||
|
import pytest |
||||
|
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 |
||||
|
from fastapi import format_sse |
||||
|
assert SSEResponse is not None |
||||
|
assert format_sse is not None |
||||
Loading…
Reference in new issue