brain-of-reese/tests/unit/test_response_budget.py

"""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<method>[A-Z]+) path=(?P<path>\S+) status=(?P<status>\d+) "
    r"bytes=(?P<bytes>\d+) enc=(?P<enc>\S+) duration_ms=(?P<ms>\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/"