"""Unit tests: ``app.core.response_budget`` — the byte-budget log line (phase 167, task 05, D5, LOCKED A5). Pins the contract of the line, the level cut, and the pass-through rules on a tiny ASGI app (the raw-ASGI interleaving harness of ``test_compression.py`` where ordering matters, ``TestClient`` where the composed stack matters): * the EXACT shape ``api method=… path=… status=… bytes=… enc=… duration_ms=…`` for a large 200 at INFO and a small 200 at DEBUG; * a 304 logs ``bytes=0`` (the revalidation win stays visible); * a non-``/api/`` path and a non-HTTP scope log NOTHING; * a stream logs ONE line with ``stream=1``, ``bytes`` = the sum of the chunks, emitted when the stream COMPLETES (never mid-flight) — and the frames are never buffered; * ``enc`` reads the WIRE encoding (``gzip`` through the real compression layer, ``identity`` for a client that does not accept it) while ``bytes`` stays the size the ROUTE produced; * an app exception is logged (finally-style, bytes as counted) and re-raised UNCHANGED; * the config knob: ``0`` logs every ``/api/*`` response at INFO, a negative value refuses startup naming the field, and the middleware sits INSIDE ``SelectiveGZipMiddleware`` so ``bytes=`` is the route's size. """ from __future__ import annotations import asyncio import gzip import json import logging import re from typing import Any import pytest from fastapi.testclient import TestClient from pydantic import ValidationError from starlette.applications import Starlette from starlette.responses import Response from starlette.routing import Route from starlette.types import Message, Receive, Scope, Send import app.core.response_budget as budget_mod import app.main as main_mod from app.config import Settings from app.core.compression import SelectiveGZipMiddleware from app.core.response_budget import ( API_PATH_PREFIX, UNKNOWN_STATUS, ResponseBudgetMiddleware, ) #: One line, fully anchored — the shape README documents and the #: ``grep ' :: api '`` workflow parses. LINE_RE = re.compile( r"^api method=(?P[A-Z]+) path=(?P\S+) status=(?P\d+) " r"bytes=(?P\d+) enc=(?P\S+) duration_ms=(?P\d+)$" ) #: ~28 KB of tree-shaped JSON. _BODY = json.dumps( { "sources": [ { "name": "src", "documents": i, "children": [{"path": f"d/{i}/{n}.md"} for n in range(20)], } for i in range(200) ] } ).encode() _SMALL_BODY = b'{"ok":true}' #: Four SSE-class frames, each comfortably above the gzip floor. _FRAMES = [f"data: {'x' * 500}\n\n".encode() for _ in range(4)] _LOGGER_NAME = "app.response_budget" def _settings(**kwargs: Any) -> Settings: """Build Settings without reading a .env file (the test_config pattern).""" kwargs.setdefault("_env_file", None) return Settings(**kwargs) # pyright: ignore[reportCallIssue] (kwarg exists at runtime) def _api_lines(caplog: pytest.LogCaptureFixture) -> list[tuple[int, str]]: """``(levelno, message)`` for this module's logger only, in order.""" return [ (record.levelno, record.getMessage()) for record in caplog.records if record.name == _LOGGER_NAME ] class _Counter(logging.Handler): """Counts emitted lines live, so an app can observe "has the line been written yet?" from INSIDE the response — the ordering proof.""" def __init__(self) -> None: super().__init__() self.count = 0 def emit(self, record: logging.LogRecord) -> None: self.count += 1 def _http_scope(path: str, method: str = "GET") -> Scope: return {"type": "http", "method": method, "path": path, "headers": []} async def _receive() -> Message: return {"type": "http.disconnect"} async def _send(message: Message) -> None: return None def _raw_json_app(body: bytes, *, status: int = 200, media_type: str = "application/json") -> Any: """The smallest ASGI app that answers one buffered body — the raw-ASGI level is where the WIRE bytes (post-compression) are observable.""" async def app(scope: Scope, receive: Receive, send: Send) -> None: await send( { "type": "http.response.start", "status": status, "headers": [(b"content-type", media_type.encode())], } ) await send({"type": "http.response.body", "body": body, "more_body": False}) return app def _raw_harness(middleware_app: Any, *, path: str, accept: str) -> list[Message]: """Drive a middleware stack at the raw ASGI level; return the exact messages that reached the client's ``send``.""" sent: list[Message] = [] async def recording_send(message: Message) -> None: sent.append(message) asyncio.run( middleware_app( { "type": "http", "method": "GET", "path": path, "headers": [(b"accept-encoding", accept.encode())], }, _receive, recording_send, ) ) return sent def _inner_app( *, body: bytes = _BODY, status: int = 200, path: str = "/api/docs/tree", media_type: str = "application/json", ) -> Starlette: async def respond(request: Any) -> Response: return Response(body, status_code=status, media_type=media_type) async def boom(request: Any) -> Response: raise RuntimeError("kaboom") return Starlette( routes=[ Route(path, respond, methods=["GET", "POST"]), Route("/api/boom", boom, methods=["GET"]), Route("/sources.html", respond, methods=["GET"]), ] ) def _client(app: Any, *, threshold: int = 1024) -> TestClient: return TestClient(ResponseBudgetMiddleware(app, threshold=threshold)) # --------------------------------------------------------------------------- # The line's exact shape, and the level the budget decides # --------------------------------------------------------------------------- def test_large_response_logs_one_info_line_in_the_exact_shape( caplog: pytest.LogCaptureFixture, ) -> None: with caplog.at_level(logging.DEBUG, logger=_LOGGER_NAME): client = _client(_inner_app(), threshold=1024) assert client.get("/api/docs/tree").status_code == 200 lines = _api_lines(caplog) assert len(lines) == 1 level, message = lines[0] assert level == logging.INFO match = LINE_RE.match(message) assert match is not None, message assert match.group("method") == "GET" assert match.group("path") == "/api/docs/tree" assert match.group("status") == "200" assert int(match.group("bytes")) == len(_BODY) assert match.group("enc") == "identity" assert int(match.group("ms")) >= 0 def test_small_response_logs_the_same_line_at_debug( caplog: pytest.LogCaptureFixture, ) -> None: with caplog.at_level(logging.DEBUG, logger=_LOGGER_NAME): client = _client(_inner_app(body=_SMALL_BODY), threshold=1024) assert client.get("/api/docs/tree").status_code == 200 lines = _api_lines(caplog) assert len(lines) == 1 level, message = lines[0] assert level == logging.DEBUG assert f"bytes={len(_SMALL_BODY)}" in message def test_budget_is_the_inclusive_cut(caplog: pytest.LogCaptureFixture) -> None: """``bytes >= budget`` is INFO, one byte less is DEBUG — the boundary the README and ``.env.example`` state.""" async def exactly(request: Any) -> Response: return Response(b"x" * 100) async def one_less(request: Any) -> Response: return Response(b"x" * 99) app = Starlette( routes=[ Route("/api/a", exactly, methods=["GET"]), Route("/api/b", one_less, methods=["GET"]), ] ) with caplog.at_level(logging.DEBUG, logger=_LOGGER_NAME): client = _client(app, threshold=100) client.get("/api/a") client.get("/api/b") assert [level for level, _ in _api_lines(caplog)] == [logging.INFO, logging.DEBUG] def test_threshold_defaults_to_the_settings_value( monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture ) -> None: """No ``threshold`` argument — the way ``create_app`` wires it — so the live ``log_large_response_bytes`` decides per response.""" monkeypatch.setattr( budget_mod, "get_settings", lambda: _settings(log_large_response_bytes=5) ) with caplog.at_level(logging.DEBUG, logger=_LOGGER_NAME): TestClient(ResponseBudgetMiddleware(_inner_app(body=_SMALL_BODY))).get( "/api/docs/tree" ) assert [level for level, _ in _api_lines(caplog)] == [logging.INFO] def test_zero_budget_logs_every_api_response_at_info( monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture ) -> None: """``BOR_LOG_LARGE_RESPONSE_BYTES=0`` is the loud mode: even an 11-byte body is INFO (the operator's one-off investigation setting).""" monkeypatch.setattr( budget_mod, "get_settings", lambda: _settings(log_large_response_bytes=0) ) with caplog.at_level(logging.DEBUG, logger=_LOGGER_NAME): TestClient(ResponseBudgetMiddleware(_inner_app(body=_SMALL_BODY))).get( "/api/docs/tree" ) # 11 bytes, budget 0 -> INFO (the same body is DEBUG at any positive # budget: see test_small_response_logs_the_same_line_at_debug). assert [level for level, _ in _api_lines(caplog)] == [logging.INFO] # --------------------------------------------------------------------------- # Statuses: a 304 logs 0 bytes, an error is counted like any other response # --------------------------------------------------------------------------- def test_304_logs_zero_bytes(caplog: pytest.LogCaptureFixture) -> None: """The revalidation win stays in the log — a bodiless 304 is a line with ``bytes=0``, not a missing line.""" with caplog.at_level(logging.DEBUG, logger=_LOGGER_NAME): client = _client(_inner_app(body=b"", status=304), threshold=1024) assert client.get("/api/docs/tree").status_code == 304 lines = _api_lines(caplog) assert len(lines) == 1 match = LINE_RE.match(lines[0][1]) assert match is not None assert match.group("status") == "304" assert match.group("bytes") == "0" assert lines[0][0] == logging.DEBUG def test_error_status_is_logged_like_any_other_response( caplog: pytest.LogCaptureFixture, ) -> None: with caplog.at_level(logging.DEBUG, logger=_LOGGER_NAME): client = _client(_inner_app(body=b'{"detail":"admin only"}', status=403)) assert client.get("/api/docs/tree").status_code == 403 match = LINE_RE.match(_api_lines(caplog)[0][1]) assert match is not None and match.group("status") == "403" # --------------------------------------------------------------------------- # What is NOT measured # --------------------------------------------------------------------------- def test_non_api_path_logs_nothing(caplog: pytest.LogCaptureFixture) -> None: with caplog.at_level(logging.DEBUG, logger=_LOGGER_NAME): # budget 0: ANY /api line would show, so silence is the whole pin. client = _client(_inner_app(), threshold=0) assert client.get("/sources.html").status_code == 200 assert _api_lines(caplog) == [] def test_non_http_scopes_pass_through_silently( caplog: pytest.LogCaptureFixture, ) -> None: """``lifespan`` / ``websocket`` scopes reach the app untouched and write no line — there is no response to measure.""" seen: list[str] = [] async def app(scope: Scope, receive: Receive, send: Send) -> None: seen.append(scope["type"]) with caplog.at_level(logging.DEBUG, logger=_LOGGER_NAME): for scope_type in ("lifespan", "websocket"): asyncio.run( ResponseBudgetMiddleware(app, threshold=0)( {"type": scope_type, "path": "/api/x"}, _receive, _send ) ) assert seen == ["lifespan", "websocket"] assert _api_lines(caplog) == [] # --------------------------------------------------------------------------- # The stream: one line, stream=1, counted total, emitted at completion # --------------------------------------------------------------------------- def test_stream_logs_once_at_completion_with_stream_flag_and_total_bytes( caplog: pytest.LogCaptureFixture, ) -> None: """The SSE contract of the line: nothing while frames are in flight, one ``stream=1`` line when the final chunk completes, ``bytes`` = the sum of the frames (NOT one frame).""" counter = _Counter() logger = logging.getLogger(_LOGGER_NAME) logger.addHandler(counter) in_flight: list[int] = [] async def sse_app(scope: Scope, receive: Receive, send: Send) -> None: await send( { "type": "http.response.start", "status": 200, "headers": [(b"content-type", b"text/event-stream")], } ) for frame in _FRAMES: await send( {"type": "http.response.body", "body": frame, "more_body": True} ) in_flight.append(counter.count) # after each frame is delivered await send({"type": "http.response.body", "body": b"", "more_body": False}) in_flight.append(counter.count) try: with caplog.at_level(logging.DEBUG, logger=_LOGGER_NAME): asyncio.run( ResponseBudgetMiddleware(sse_app, threshold=1)( _http_scope("/api/chat", "POST"), _receive, _send ) ) finally: logger.removeHandler(counter) # Every in-flight probe except the last saw zero lines; the completion # is what emitted. assert in_flight == [0] * len(_FRAMES) + [1] lines = _api_lines(caplog) assert len(lines) == 1 message = lines[0][1] assert message.endswith(" stream=1") match = LINE_RE.match(message.removesuffix(" stream=1")) assert match is not None assert match.group("method") == "POST" assert match.group("path") == "/api/chat" assert int(match.group("bytes")) == sum(len(frame) for frame in _FRAMES) def test_stream_frames_are_never_buffered() -> None: """Zero-buffering proof (the interleaving idiom of ``test_compression.py``): every frame reaches the client's ``send`` before the app produces the next one — a ``BaseHTTPMiddleware`` here would emit all ``app:*`` events first.""" events: list[str] = [] async def sse_app(scope: Scope, receive: Receive, send: Send) -> None: await send( { "type": "http.response.start", "status": 200, "headers": [(b"content-type", b"text/event-stream")], } ) for i, frame in enumerate(_FRAMES): events.append(f"app:frame-{i}") await send( {"type": "http.response.body", "body": frame, "more_body": True} ) await send({"type": "http.response.body", "body": b"", "more_body": False}) async def recording_send(message: Message) -> None: events.append(f"client:{message['type']}") asyncio.run( ResponseBudgetMiddleware(sse_app, threshold=0)( _http_scope("/api/chat", "POST"), _receive, recording_send ) ) expected = ["client:http.response.start"] for i in range(len(_FRAMES)): expected += [f"app:frame-{i}", "client:http.response.body"] expected += ["client:http.response.body"] assert events == expected # --------------------------------------------------------------------------- # enc= reads the wire, bytes= stays the route's size # --------------------------------------------------------------------------- def _stack(*, threshold: int = 1024) -> TestClient: """response-budget INSIDE the selective-gzip layer — the order ``create_app`` builds.""" return TestClient( SelectiveGZipMiddleware( ResponseBudgetMiddleware(_inner_app(), threshold=threshold), minimum_size=860, ) ) def test_enc_gzip_is_reported_through_the_compression_layer( caplog: pytest.LogCaptureFixture, ) -> None: """The two halves of the line, measured against the wire: the client receives a gzipped body, the log says ``enc=gzip`` and still counts the ROUTE's uncompressed bytes (the position inside the compressor is what buys that).""" with caplog.at_level(logging.DEBUG, logger=_LOGGER_NAME): sent = _raw_harness( SelectiveGZipMiddleware( ResponseBudgetMiddleware(_raw_json_app(_BODY), threshold=1024), minimum_size=860, ), path="/api/docs/tree", accept="gzip", ) wire = b"".join(m.get("body", b"") for m in sent if m["type"] == "http.response.body") assert gzip.decompress(wire) == _BODY # the transfer really was gzip match = LINE_RE.match(_api_lines(caplog)[-1][1]) assert match is not None assert match.group("enc") == "gzip" assert int(match.group("bytes")) == len(_BODY) > len(wire) def test_enc_identity_when_the_client_does_not_accept_gzip( caplog: pytest.LogCaptureFixture, ) -> None: with caplog.at_level(logging.DEBUG, logger=_LOGGER_NAME): response = _stack().get( "/api/docs/tree", headers={"Accept-Encoding": "identity"} ) assert "content-encoding" not in response.headers match = LINE_RE.match(_api_lines(caplog)[-1][1]) assert match is not None and match.group("enc") == "identity" assert int(match.group("bytes")) == len(_BODY) def test_enc_reports_a_pre_encoded_response(caplog: pytest.LogCaptureFixture) -> None: """A route that encodes its own body (the fronting-proxy case) is reported as such — the header on the wire is the answer.""" async def encoded(request: Any) -> Response: return Response( gzip.compress(_BODY), media_type="application/json", headers={"Content-Encoding": "gzip"}, ) app = Starlette(routes=[Route("/api/encoded", encoded, methods=["GET"])]) with caplog.at_level(logging.DEBUG, logger=_LOGGER_NAME): TestClient(ResponseBudgetMiddleware(app, threshold=1024)).get("/api/encoded") match = LINE_RE.match(_api_lines(caplog)[-1][1]) assert match is not None and match.group("enc") == "gzip" # --------------------------------------------------------------------------- # The app raising: logged, and re-raised unchanged # --------------------------------------------------------------------------- def test_app_exception_logs_and_reraises_unchanged( caplog: pytest.LogCaptureFixture, ) -> None: """An app that dies before it can answer still writes one line (nothing was counted, status 500 — what the outer error layer will answer) and the exception propagates UNCHANGED.""" async def dying_app(scope: Scope, receive: Receive, send: Send) -> None: raise RuntimeError("kaboom") with ( caplog.at_level(logging.DEBUG, logger=_LOGGER_NAME), pytest.raises(RuntimeError, match="kaboom"), ): asyncio.run( ResponseBudgetMiddleware(dying_app, threshold=0)( _http_scope("/api/docs/tree"), _receive, _send ) ) lines = _api_lines(caplog) assert len(lines) == 1 match = LINE_RE.match(lines[0][1]) assert match is not None assert match.group("status") == str(UNKNOWN_STATUS) assert match.group("bytes") == "0" def test_a_500_produced_by_the_framework_is_counted_like_a_response( caplog: pytest.LogCaptureFixture, ) -> None: """Through a real app the unhandled exception becomes a 500 WITH a body (Starlette's ``Internal Server Error`` text): the line then reports that status and those bytes — one line either way, never zero.""" with ( caplog.at_level(logging.DEBUG, logger=_LOGGER_NAME), pytest.raises(RuntimeError, match="kaboom"), ): _client(_inner_app(), threshold=0).get("/api/boom") lines = _api_lines(caplog) assert len(lines) == 1 match = LINE_RE.match(lines[0][1]) assert match is not None assert match.group("status") == "500" assert int(match.group("bytes")) == len(b"Internal Server Error") def test_exception_mid_stream_logs_the_bytes_counted_so_far( caplog: pytest.LogCaptureFixture, ) -> None: """A stream that dies after two frames still writes ONE line with the bytes it actually got out (finally-style) — and the exception is not swallowed.""" async def dying_stream_app(scope: Scope, receive: Receive, send: Send) -> None: await send( { "type": "http.response.start", "status": 200, "headers": [(b"content-type", b"text/event-stream")], } ) await send({"type": "http.response.body", "body": _FRAMES[0], "more_body": True}) await send({"type": "http.response.body", "body": _FRAMES[1], "more_body": True}) raise RuntimeError("stream died") with ( caplog.at_level(logging.DEBUG, logger=_LOGGER_NAME), pytest.raises(RuntimeError, match="stream died"), ): asyncio.run( ResponseBudgetMiddleware(dying_stream_app, threshold=0)( _http_scope("/api/chat", "POST"), _receive, _send ) ) lines = _api_lines(caplog) assert len(lines) == 1 message = lines[0][1] assert message.endswith(" stream=1") assert f"bytes={len(_FRAMES[0]) + len(_FRAMES[1])}" in message def test_a_cancelled_stream_still_logs_one_line_with_what_got_out( caplog: pytest.LogCaptureFixture, ) -> None: """The client-walked-away path (the live audit hit it: an SSE turn cut mid-answer): the app unwinds with ``CancelledError`` and no completing body, the ``finally`` writes ONE line with the bytes that actually got out, and the cancellation propagates (the middleware never swallows it).""" async def cancelled_stream_app(scope: Scope, receive: Receive, send: Send) -> None: await send( { "type": "http.response.start", "status": 200, "headers": [(b"content-type", b"text/event-stream")], } ) await send({"type": "http.response.body", "body": _FRAMES[0], "more_body": True}) raise asyncio.CancelledError # what the server does on a client disconnect with ( caplog.at_level(logging.DEBUG, logger=_LOGGER_NAME), pytest.raises(asyncio.CancelledError), ): asyncio.run( ResponseBudgetMiddleware(cancelled_stream_app, threshold=0)( _http_scope("/api/chat", "POST"), _receive, _send ) ) lines = _api_lines(caplog) assert len(lines) == 1 assert lines[0][1].endswith(" stream=1") assert f"bytes={len(_FRAMES[0])}" in lines[0][1] # --------------------------------------------------------------------------- # Config + wiring # --------------------------------------------------------------------------- def test_log_large_response_bytes_defaults_to_64_kb() -> None: assert _settings().log_large_response_bytes == 65536 def test_log_large_response_bytes_env_override( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setenv("BOR_LOG_LARGE_RESPONSE_BYTES", "4096") assert _settings().log_large_response_bytes == 4096 monkeypatch.setenv("BOR_LOG_LARGE_RESPONSE_BYTES", "0") # loud mode is legal assert _settings().log_large_response_bytes == 0 def test_negative_budget_refuses_startup_naming_the_field() -> None: with pytest.raises(ValidationError, match="log_large_response_bytes must be >= 0"): _settings(log_large_response_bytes=-1) def test_middleware_sits_inside_gzip_outside_caching( monkeypatch: pytest.MonkeyPatch, ) -> None: """The documented stack (README + ``app/main.py``): SecurityHeaders outermost → gzip → response-budget → caching → session. That position is what makes ``bytes=`` the ROUTE's size while ``enc=`` still names the wire encoding.""" monkeypatch.setattr(main_mod.settings, "compress_min_bytes", 860) names = [ getattr(mw.cls, "__name__", "") for mw in main_mod.create_app().user_middleware ] assert names == [ "SecurityHeadersMiddleware", "SelectiveGZipMiddleware", "ResponseBudgetMiddleware", "CachingMiddleware", "SessionMiddleware", ] def test_budget_is_wired_even_with_compression_off( monkeypatch: pytest.MonkeyPatch, ) -> None: """``BOR_COMPRESS_MIN_BYTES=0`` removes gzip only — the budget line is not optional (LOCKED A5: measure everything, reject nothing), so the measurement survives the compression kill switch.""" monkeypatch.setattr(main_mod.settings, "compress_min_bytes", 0) names = [ getattr(mw.cls, "__name__", "") for mw in main_mod.create_app().user_middleware ] assert "SelectiveGZipMiddleware" not in names assert "ResponseBudgetMiddleware" in names def test_api_path_prefix_constant() -> None: assert API_PATH_PREFIX == "/api/"