diff --git a/fastapi/routing.py b/fastapi/routing.py index c442b122b..8a4819b29 100644 --- a/fastapi/routing.py +++ b/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: diff --git a/tests/test_stream_status_code.py b/tests/test_stream_status_code.py new file mode 100644 index 000000000..76b2bf4dc --- /dev/null +++ b/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)