Browse Source

🐛 Fix status_code being ignored for SSE and JSONL streaming endpoints

pull/15937/head
Saurabh 3 weeks ago
parent
commit
8c2b6c873e
  1. 10
      fastapi/routing.py
  2. 138
      tests/test_stream_status_code.py

10
fastapi/routing.py

@ -624,10 +624,13 @@ def get_request_handler(
_sse_with_checkpoints(sse_receive_stream)
)
response_args = _build_response_args(
status_code=status_code, solved_result=solved_result
)
response = StreamingResponse(
sse_stream_content,
media_type="text/event-stream",
background=solved_result.background_tasks,
**response_args,
)
response.headers["Cache-Control"] = "no-cache"
# For Nginx proxies to not buffer server sent events
@ -660,10 +663,13 @@ def get_request_handler(
jsonl_stream_content = _sync_stream_jsonl()
response_args = _build_response_args(
status_code=status_code, solved_result=solved_result
)
response = StreamingResponse(
jsonl_stream_content,
media_type="application/jsonl",
background=solved_result.background_tasks,
**response_args,
)
response.headers.raw.extend(solved_result.response.headers.raw)
elif dependant.is_async_gen_callable or dependant.is_gen_callable:

138
tests/test_stream_status_code.py

@ -0,0 +1,138 @@
import json
from collections.abc import AsyncIterable
from typing import Any, cast
from fastapi import Depends, FastAPI, Response
from fastapi.responses import EventSourceResponse, StreamingResponse
from fastapi.testclient import TestClient
app = FastAPI()
def set_accepted(response: Response) -> None:
response.status_code = 202
@app.post("/sse-created", response_class=EventSourceResponse, status_code=201)
async def sse_created() -> AsyncIterable[dict[str, str]]:
yield {"message": "created"}
@app.post("/jsonl-created", status_code=201)
async def jsonl_created() -> AsyncIterable[dict[str, str]]:
yield {"message": "created"}
@app.get("/sse-dependency", response_class=EventSourceResponse)
async def sse_dependency(
accepted: None = Depends(set_accepted),
) -> AsyncIterable[dict[str, str]]:
yield {"message": "accepted"}
@app.get("/jsonl-dependency")
async def jsonl_dependency(
accepted: None = Depends(set_accepted),
) -> AsyncIterable[dict[str, str]]:
yield {"message": "accepted"}
@app.get("/raw-dependency", response_class=StreamingResponse)
async def raw_dependency(
accepted: None = Depends(set_accepted),
) -> AsyncIterable[str]:
yield "accepted"
@app.get("/sse-default", response_class=EventSourceResponse)
async def sse_default() -> AsyncIterable[dict[str, str]]:
yield {"message": "ok"}
@app.get("/jsonl-default")
async def jsonl_default() -> AsyncIterable[dict[str, str]]:
yield {"message": "ok"}
client = TestClient(app)
def get_sse_data(response_text: str) -> list[dict[str, str]]:
data_lines = [
line.removeprefix("data: ")
for line in response_text.splitlines()
if line.startswith("data: ")
]
return [json.loads(line) for line in data_lines]
def get_jsonl_data(response_text: str) -> list[dict[str, str]]:
return [json.loads(line) for line in response_text.splitlines()]
def test_sse_stream_honors_declared_status_code() -> None:
response = client.post("/sse-created")
assert response.status_code == 201
assert response.headers["content-type"] == "text/event-stream; charset=utf-8"
assert get_sse_data(response.text) == [{"message": "created"}]
def test_jsonl_stream_honors_declared_status_code() -> None:
response = client.post("/jsonl-created")
assert response.status_code == 201
assert response.headers["content-type"] == "application/jsonl"
assert get_jsonl_data(response.text) == [{"message": "created"}]
def test_sse_stream_honors_dependency_status_code() -> None:
response = client.get("/sse-dependency")
assert response.status_code == 202
assert get_sse_data(response.text) == [{"message": "accepted"}]
def test_jsonl_stream_honors_dependency_status_code() -> None:
response = client.get("/jsonl-dependency")
assert response.status_code == 202
assert get_jsonl_data(response.text) == [{"message": "accepted"}]
def test_raw_stream_still_honors_dependency_status_code() -> None:
response = client.get("/raw-dependency")
assert response.status_code == 202
assert response.text == "accepted"
def test_sse_stream_default_status_code_stays_200() -> None:
response = client.get("/sse-default")
assert response.status_code == 200
assert get_sse_data(response.text) == [{"message": "ok"}]
def test_jsonl_stream_default_status_code_stays_200() -> None:
response = client.get("/jsonl-default")
assert response.status_code == 200
assert get_jsonl_data(response.text) == [{"message": "ok"}]
def test_stream_status_codes_match_openapi() -> None:
schema = client.get("/openapi.json").json()
assert response_status_codes(schema, "/sse-created", "post") == ["201"]
assert response_status_codes(schema, "/jsonl-created", "post") == ["201"]
assert response_status_codes(schema, "/sse-default", "get") == ["200"]
assert response_status_codes(schema, "/jsonl-default", "get") == ["200"]
def response_status_codes(schema: dict[str, Any], path: str, method: str) -> list[str]:
paths = cast(dict[str, Any], schema["paths"])
route = cast(dict[str, Any], paths[path])
operation = cast(dict[str, Any], route[method])
responses = cast(dict[str, Any], operation["responses"])
return sorted(responses)
Loading…
Cancel
Save