Browse Source

Update tests

pull/15937/head
Yurii Motov 1 week ago
parent
commit
8938eda5e9
  1. 500
      tests/test_stream_status_code.py

500
tests/test_stream_status_code.py

@ -1,10 +1,50 @@
import json
from collections.abc import AsyncIterable from collections.abc import AsyncIterable
from typing import Any, cast
import pytest
from fastapi import Depends, FastAPI, Response from fastapi import Depends, FastAPI, Response
from fastapi.responses import EventSourceResponse, StreamingResponse from fastapi.responses import EventSourceResponse, StreamingResponse
from fastapi.testclient import TestClient from fastapi.testclient import TestClient
from inline_snapshot import snapshot
SSE_RESPONSE = {
"description": "Successful Response",
"content": {
"text/event-stream": {
"itemSchema": {
"type": "object",
"properties": {
"data": {
"type": "string",
"contentMediaType": "application/json",
"contentSchema": {
"type": "object",
"additionalProperties": {"type": "string"},
"title": "SSE stream item",
},
},
"event": {"type": "string"},
"id": {"type": "string"},
"retry": {"type": "integer", "minimum": 0},
},
"required": ["data"],
}
}
},
}
JSONL_RESPONSE = {
"description": "Successful Response",
"content": {
"application/jsonl": {
"itemSchema": {
"type": "object",
"additionalProperties": {"type": "string"},
"title": "JSONL stream item",
}
}
},
}
app = FastAPI() app = FastAPI()
@ -13,53 +53,80 @@ def set_accepted(response: Response) -> None:
response.status_code = 202 response.status_code = 202
@app.post("/sse-created", response_class=EventSourceResponse, status_code=201) @app.post("/sse", response_class=EventSourceResponse, status_code=201)
async def sse_created() -> AsyncIterable[dict[str, str]]: async def sse() -> AsyncIterable[dict[str, str]]:
yield {"message": "created"} yield {"message": "created"}
@app.post("/jsonl-created", status_code=201) @app.post("/jsonl", status_code=201)
async def jsonl_created() -> AsyncIterable[dict[str, str]]: async def jsonl() -> AsyncIterable[dict[str, str]]:
yield {"message": "created"} yield {"message": "created"}
@app.get("/sse-dependency", response_class=EventSourceResponse) @app.post("/raw", response_class=StreamingResponse, status_code=201)
async def raw() -> AsyncIterable[str]:
yield "accepted"
@app.post(
"/sse-dependency",
response_class=EventSourceResponse,
responses={202: SSE_RESPONSE},
)
async def sse_dependency( async def sse_dependency(
accepted: None = Depends(set_accepted), accepted: None = Depends(set_accepted),
) -> AsyncIterable[dict[str, str]]: ) -> AsyncIterable[dict[str, str]]:
yield {"message": "accepted"} yield {"message": "accepted"}
@app.get("/jsonl-dependency") @app.post("/jsonl-dependency", responses={202: JSONL_RESPONSE})
async def jsonl_dependency( async def jsonl_dependency(
accepted: None = Depends(set_accepted), accepted: None = Depends(set_accepted),
) -> AsyncIterable[dict[str, str]]: ) -> AsyncIterable[dict[str, str]]:
yield {"message": "accepted"} yield {"message": "accepted"}
@app.get("/raw-dependency", response_class=StreamingResponse) @app.post(
"/raw-dependency",
response_class=StreamingResponse,
responses={202: {"description": "Accepted"}},
)
async def raw_dependency( async def raw_dependency(
accepted: None = Depends(set_accepted), accepted: None = Depends(set_accepted),
) -> AsyncIterable[str]: ) -> AsyncIterable[str]:
yield "accepted" yield "accepted"
@app.post("/sse-created-override", response_class=EventSourceResponse, status_code=201) @app.post(
async def sse_created_override( "/sse-dependency-override",
response_class=EventSourceResponse,
status_code=201,
responses={202: SSE_RESPONSE},
)
async def sse_dependency_override(
accepted: None = Depends(set_accepted), accepted: None = Depends(set_accepted),
) -> AsyncIterable[dict[str, str]]: ) -> AsyncIterable[dict[str, str]]:
yield {"message": "overridden"} yield {"message": "overridden"}
@app.post("/jsonl-created-override", status_code=201) @app.post(
async def jsonl_created_override( "/jsonl-dependency-override",
status_code=201,
responses={202: JSONL_RESPONSE},
)
async def jsonl_dependency_override(
accepted: None = Depends(set_accepted), accepted: None = Depends(set_accepted),
) -> AsyncIterable[dict[str, str]]: ) -> AsyncIterable[dict[str, str]]:
yield {"message": "overridden"} yield {"message": "overridden"}
@app.post("/raw-created-override", response_class=StreamingResponse, status_code=201) @app.post(
async def raw_created_override( "/raw-dependency-override",
response_class=StreamingResponse,
status_code=201,
responses={202: {"description": "Accepted"}},
)
async def raw_dependency_override(
accepted: None = Depends(set_accepted), accepted: None = Depends(set_accepted),
) -> AsyncIterable[str]: ) -> AsyncIterable[str]:
yield "overridden" yield "overridden"
@ -68,91 +135,318 @@ async def raw_created_override(
client = TestClient(app) client = TestClient(app)
def get_sse_data(response_text: str) -> list[dict[str, str]]: @pytest.mark.parametrize(
data_lines = [ "path,expected_status_code",
line.removeprefix("data: ") [
for line in response_text.splitlines() ("/sse", 201),
if line.startswith("data: ") ("/jsonl", 201),
] ("/raw", 201),
return [json.loads(line) for line in data_lines] ("/sse-dependency", 202),
("/jsonl-dependency", 202),
("/raw-dependency", 202),
def get_jsonl_data(response_text: str) -> list[dict[str, str]]: ("/sse-dependency-override", 202),
return [json.loads(line) for line in response_text.splitlines()] ("/jsonl-dependency-override", 202),
("/raw-dependency-override", 202),
],
def test_sse_stream_honors_declared_status_code() -> None: )
response = client.post("/sse-created") def test_status_code(path: str, expected_status_code: int) -> None:
response = client.post(path)
assert response.status_code == 201 assert response.status_code == expected_status_code
assert response.headers["content-type"] == "text/event-stream; charset=utf-8"
assert get_sse_data(response.text) == [{"message": "created"}]
def test_openapi() -> None:
openapi = app.openapi()
def test_jsonl_stream_honors_declared_status_code() -> None:
response = client.post("/jsonl-created") assert openapi == snapshot(
{
assert response.status_code == 201 "openapi": "3.1.0",
assert response.headers["content-type"] == "application/jsonl" "info": {"title": "FastAPI", "version": "0.1.0"},
assert get_jsonl_data(response.text) == [{"message": "created"}] "paths": {
"/sse": {
"post": {
def test_sse_stream_honors_dependency_status_code() -> None: "summary": "Sse",
response = client.get("/sse-dependency") "operationId": "sse_sse_post",
"responses": {
assert response.status_code == 202 "201": {
assert get_sse_data(response.text) == [{"message": "accepted"}] "description": "Successful Response",
"content": {
"text/event-stream": {
def test_jsonl_stream_honors_dependency_status_code() -> None: "itemSchema": {
response = client.get("/jsonl-dependency") "type": "object",
"properties": {
assert response.status_code == 202 "data": {
assert get_jsonl_data(response.text) == [{"message": "accepted"}] "type": "string",
"contentMediaType": "application/json",
"contentSchema": {
def test_raw_stream_still_honors_dependency_status_code() -> None: "type": "object",
response = client.get("/raw-dependency") "additionalProperties": {
"type": "string"
assert response.status_code == 202 },
assert response.text == "accepted" "title": "Streamitem Sse Sse Post",
},
},
def test_sse_stream_dependency_overrides_declared_status_code() -> None: "event": {"type": "string"},
response = client.post("/sse-created-override") "id": {"type": "string"},
"retry": {
assert response.status_code == 202 "type": "integer",
assert get_sse_data(response.text) == [{"message": "overridden"}] "minimum": 0,
},
},
def test_jsonl_stream_dependency_overrides_declared_status_code() -> None: "required": ["data"],
response = client.post("/jsonl-created-override") }
}
assert response.status_code == 202 },
assert get_jsonl_data(response.text) == [{"message": "overridden"}] }
},
}
def test_raw_stream_dependency_overrides_declared_status_code() -> None: },
response = client.post("/raw-created-override") "/jsonl": {
"post": {
assert response.status_code == 202 "summary": "Jsonl",
assert response.text == "overridden" "operationId": "jsonl_jsonl_post",
"responses": {
"201": {
def test_stream_status_codes_match_openapi() -> None: "description": "Successful Response",
schema = client.get("/openapi.json").json() "content": {
"application/jsonl": {
assert response_status_codes(schema, "/sse-created", "post") == ["201"] "itemSchema": {
assert response_status_codes(schema, "/jsonl-created", "post") == ["201"] "type": "object",
"additionalProperties": {"type": "string"},
assert response_status_codes(schema, "/sse-created-override", "post") == ["201"] "title": "Streamitem Jsonl Jsonl Post",
assert response_status_codes(schema, "/jsonl-created-override", "post") == ["201"] }
assert response_status_codes(schema, "/raw-created-override", "post") == ["201"] }
},
}
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]) "/raw": {
responses = cast(dict[str, Any], operation["responses"]) "post": {
return sorted(responses) "summary": "Raw",
"operationId": "raw_raw_post",
"responses": {"201": {"description": "Successful Response"}},
}
},
"/sse-dependency": {
"post": {
"summary": "Sse Dependency",
"operationId": "sse_dependency_sse_dependency_post",
"responses": {
"200": {
"description": "Successful Response",
"content": {
"text/event-stream": {
"itemSchema": {
"type": "object",
"properties": {
"data": {
"type": "string",
"contentMediaType": "application/json",
"contentSchema": {
"type": "object",
"additionalProperties": {
"type": "string"
},
"title": "Streamitem Sse Dependency Sse Dependency Post",
},
},
"event": {"type": "string"},
"id": {"type": "string"},
"retry": {
"type": "integer",
"minimum": 0,
},
},
"required": ["data"],
}
}
},
},
"202": {
"description": "Successful Response",
"content": {
"text/event-stream": {
"itemSchema": {
"type": "object",
"properties": {
"data": {
"type": "string",
"contentMediaType": "application/json",
"contentSchema": {
"type": "object",
"additionalProperties": {
"type": "string"
},
"title": "SSE stream item",
},
},
"event": {"type": "string"},
"id": {"type": "string"},
"retry": {
"type": "integer",
"minimum": 0,
},
},
"required": ["data"],
}
}
},
},
},
}
},
"/jsonl-dependency": {
"post": {
"summary": "Jsonl Dependency",
"operationId": "jsonl_dependency_jsonl_dependency_post",
"responses": {
"200": {
"description": "Successful Response",
"content": {
"application/jsonl": {
"itemSchema": {
"type": "object",
"additionalProperties": {"type": "string"},
"title": "Streamitem Jsonl Dependency Jsonl Dependency Post",
}
}
},
},
"202": {
"description": "Successful Response",
"content": {
"application/jsonl": {
"itemSchema": {
"type": "object",
"additionalProperties": {"type": "string"},
"title": "JSONL stream item",
}
}
},
},
},
}
},
"/raw-dependency": {
"post": {
"summary": "Raw Dependency",
"operationId": "raw_dependency_raw_dependency_post",
"responses": {
"200": {"description": "Successful Response"},
"202": {"description": "Accepted"},
},
}
},
"/sse-dependency-override": {
"post": {
"summary": "Sse Dependency Override",
"operationId": "sse_dependency_override_sse_dependency_override_post",
"responses": {
"201": {
"description": "Successful Response",
"content": {
"text/event-stream": {
"itemSchema": {
"type": "object",
"properties": {
"data": {
"type": "string",
"contentMediaType": "application/json",
"contentSchema": {
"type": "object",
"additionalProperties": {
"type": "string"
},
"title": "Streamitem Sse Dependency Override Sse Dependency Override Post",
},
},
"event": {"type": "string"},
"id": {"type": "string"},
"retry": {
"type": "integer",
"minimum": 0,
},
},
"required": ["data"],
}
}
},
},
"202": {
"description": "Successful Response",
"content": {
"text/event-stream": {
"itemSchema": {
"type": "object",
"properties": {
"data": {
"type": "string",
"contentMediaType": "application/json",
"contentSchema": {
"type": "object",
"additionalProperties": {
"type": "string"
},
"title": "SSE stream item",
},
},
"event": {"type": "string"},
"id": {"type": "string"},
"retry": {
"type": "integer",
"minimum": 0,
},
},
"required": ["data"],
}
}
},
},
},
}
},
"/jsonl-dependency-override": {
"post": {
"summary": "Jsonl Dependency Override",
"operationId": "jsonl_dependency_override_jsonl_dependency_override_post",
"responses": {
"201": {
"description": "Successful Response",
"content": {
"application/jsonl": {
"itemSchema": {
"type": "object",
"additionalProperties": {"type": "string"},
"title": "Streamitem Jsonl Dependency Override Jsonl Dependency Override Post",
}
}
},
},
"202": {
"description": "Successful Response",
"content": {
"application/jsonl": {
"itemSchema": {
"type": "object",
"additionalProperties": {"type": "string"},
"title": "JSONL stream item",
}
}
},
},
},
}
},
"/raw-dependency-override": {
"post": {
"summary": "Raw Dependency Override",
"operationId": "raw_dependency_override_raw_dependency_override_post",
"responses": {
"201": {"description": "Successful Response"},
"202": {"description": "Accepted"},
},
}
},
},
}
)

Loading…
Cancel
Save