Aether/tests/unit/test_api_request_context.py

from __future__ import annotations

import gzip
import json

import pytest
from fastapi import HTTPException
from starlette.requests import Request

from src.api.base.context import ApiRequestContext


def _build_request(headers: dict[str, str] | None = None) -> Request:
    return _build_request_with_body(b"", headers=headers)


def _build_request_with_body(
    body: bytes,
    headers: dict[str, str] | None = None,
) -> Request:
    header_items = [
        (str(key).encode("latin-1"), str(value).encode("latin-1"))
        for key, value in (headers or {}).items()
    ]
    scope = {
        "type": "http",
        "http_version": "1.1",
        "method": "POST",
        "scheme": "http",
        "path": "/v1/messages",
        "raw_path": b"/v1/messages",
        "query_string": b"",
        "headers": header_items,
        "client": ("127.0.0.1", 12345),
        "server": ("testserver", 80),
    }

    received = False

    async def receive() -> dict[str, object]:
        nonlocal received
        if received:
            return {"type": "http.request", "body": b"", "more_body": False}
        received = True
        return {"type": "http.request", "body": body, "more_body": False}

    request = Request(scope, receive)
    request.state.perf_metrics = {}
    return request


def _build_context(raw_body: bytes, headers: dict[str, str] | None = None) -> ApiRequestContext:
    request = _build_request(headers=headers)
    return ApiRequestContext(
        request=request,
        db=None,  # type: ignore[arg-type]
        user=None,
        api_key=None,
        request_id="req_test",
        start_time=0.0,
        request_method="POST",
        request_path="/v1/messages",
        client_ip="127.0.0.1",
        user_agent="pytest",
        original_headers=headers or {},
        query_params={},
        raw_body=raw_body,
    )


class TestApiRequestContextEnsureJsonBody:
    def test_build_prefers_request_state_request_id_over_trace_header(self) -> None:
        request = _build_request(headers={"x-trace-id": "trace-frontdoor-123"})
        request.state.request_id = "state-rid-001"

        context = ApiRequestContext.build(
            request=request,
            db=None,  # type: ignore[arg-type]
            user=None,
            api_key=None,
            raw_body=b"{}",
        )

        assert context.request_id == "state-rid-001"
        assert request.state.request_id == "state-rid-001"

    def test_build_prefers_trace_header_for_request_id(self) -> None:
        request = _build_request(headers={"x-trace-id": "trace-frontdoor-123"})

        context = ApiRequestContext.build(
            request=request,
            db=None,  # type: ignore[arg-type]
            user=None,
            api_key=None,
            raw_body=b"{}",
        )

        assert context.request_id == "trace-frontdoor-123"
        assert request.state.request_id == "trace-frontdoor-123"

    def test_build_snapshots_request_method_path_and_path_params(self) -> None:
        request = _build_request(headers={"x-trace-id": "trace-frontdoor-123"})
        request.scope["method"] = "GET"
        request.scope["path"] = "/v1beta/models/gemini-2.5-pro:generateContent"
        request.scope["raw_path"] = b"/v1beta/models/gemini-2.5-pro:generateContent"
        request.scope["path_params"] = {"model": "gemini-2.5-pro"}

        context = ApiRequestContext.build(
            request=request,
            db=None,  # type: ignore[arg-type]
            user=None,
            api_key=None,
            raw_body=b"{}",
        )

        assert context.request_method == "GET"
        assert context.request_path == "/v1beta/models/gemini-2.5-pro:generateContent"
        assert context.path_params == {"model": "gemini-2.5-pro"}

    def test_build_snapshots_request_runtime_state(self) -> None:
        request = _build_request(headers={"x-trace-id": "trace-frontdoor-123"})
        request.state.prefetched_balance_remaining = "12.5"
        request.state.gateway_execution_path = "public_proxy_after_executor_miss"
        request.state.rate_limit_scope = "user"
        request.state.tx_committed_by_route = True

        context = ApiRequestContext.build(
            request=request,
            db=None,  # type: ignore[arg-type]
            user=None,
            api_key=None,
            raw_body=b"{}",
        )

        assert context.prefetched_balance_remaining == 12.5
        assert context.gateway_execution_path == "public_proxy_after_executor_miss"
        assert context.rate_limit_scope == "user"
        assert context.tx_committed_by_route is True

    def test_decompresses_gzip_body(self) -> None:
        payload = {"message": "hello", "count": 2}
        raw_body = gzip.compress(json.dumps(payload).encode("utf-8"))
        context = _build_context(raw_body, headers={"content-encoding": "gzip"})

        result = context.ensure_json_body()

        assert result == payload

    def test_rejects_invalid_gzip_body(self) -> None:
        context = _build_context(b"not-gzip-body", headers={"content-encoding": "gzip"})

        with pytest.raises(HTTPException) as exc_info:
            context.ensure_json_body()

        assert exc_info.value.status_code == 400
        assert exc_info.value.detail == "gzip 请求体解压失败"

    def test_build_records_client_encoding_preferences(self) -> None:
        request = _build_request(
            headers={
                "content-type": "application/json",
                "content-encoding": "gzip",
                "accept-encoding": "gzip, deflate",
            }
        )
        context = ApiRequestContext.build(
            request=request,
            db=None,  # type: ignore[arg-type]
            user=None,
            api_key=None,
            raw_body=b"{}",
        )

        assert context.client_content_encoding == "gzip"
        assert context.client_accept_encoding == "gzip, deflate"
        assert context.request_content_type == "application/json"

    def test_build_records_perf_only_when_payload_not_empty(self) -> None:
        request = _build_request(headers={"x-trace-id": "trace-frontdoor-123"})
        request.state.perf_metrics = {}
        context = ApiRequestContext.build(
            request=request,
            db=None,  # type: ignore[arg-type]
            user=None,
            api_key=None,
            raw_body=b"{}",
        )
        assert "perf" not in context.extra

        request_with_perf = _build_request(headers={"x-trace-id": "trace-frontdoor-456"})
        request_with_perf.state.perf_metrics = {"pipeline": {"auth_ms": 3}}
        context_with_perf = ApiRequestContext.build(
            request=request_with_perf,
            db=None,  # type: ignore[arg-type]
            user=None,
            api_key=None,
            raw_body=b"{}",
        )
        assert context_with_perf.extra["perf"] == {"pipeline": {"auth_ms": 3}}
        assert context_with_perf.perf_metrics == {"pipeline": {"auth_ms": 3}}

    @pytest.mark.asyncio
    async def test_ensure_json_body_async_loads_body_lazily(self) -> None:
        payload = {"message": "hello", "count": 2}
        request = _build_request_with_body(json.dumps(payload).encode("utf-8"))
        context = ApiRequestContext.build(
            request=request,
            db=None,  # type: ignore[arg-type]
            user=None,
            api_key=None,
            raw_body=None,
        )

        assert context.raw_body is None

        result = await context.ensure_json_body_async()

        assert result == payload
        assert context.raw_body == json.dumps(payload).encode("utf-8")