chopratejas--headroom/tests/test_auth_mode.py

"""Python parity tests for ``headroom.proxy.auth_mode.classify_auth_mode``.

Mirrors the Rust matrix in ``crates/headroom-core/tests/auth_mode.rs``
byte-for-byte. The two implementations MUST agree on every header set
covered here. Any divergence is a bug — we catch it at PR review by
running both suites side by side.
"""

from __future__ import annotations

import time

import pytest

from headroom.proxy.auth_mode import (
    SUBSCRIPTION_UA_PREFIXES,
    AuthMode,
    classify_auth_mode,
    classify_client,
)


def _h(**pairs: str) -> dict[str, str]:
    """Build a plain-dict header set in one expression."""
    return dict(pairs)


# ── Required matrix ──────────────────────────────────────────────


def test_api_key_classified_payg() -> None:
    """Anthropic PAYG: ``Authorization: Bearer sk-ant-api03-...``."""
    headers = {"authorization": "Bearer sk-ant-api03-abc123def456"}
    assert classify_auth_mode(headers) is AuthMode.PAYG


def test_oauth_jwt_classified_oauth() -> None:
    """Codex / Cursor OAuth bearer: classic 3-segment JWT."""
    jwt = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0In0.signaturepart"
    headers = {"authorization": f"Bearer {jwt}"}
    assert classify_auth_mode(headers) is AuthMode.OAUTH


def test_oauth_sk_ant_oat_classified_oauth() -> None:
    """Claude Pro / Max OAuth: ``Bearer sk-ant-oat-...``."""
    headers = {"authorization": "Bearer sk-ant-oat-01-abc123def456"}
    assert classify_auth_mode(headers) is AuthMode.OAUTH


def test_oauth_real_sk_ant_oat01_classified_oauth() -> None:
    """Real Anthropic OAuth access tokens are ``sk-ant-oat01-...`` (a version
    number, no dash after ``oat``). These must classify as OAUTH — matching on
    ``sk-ant-oat-`` missed them and let them fall through to PAYG, enabling
    aggressive lossy compression on subscription-bound requests."""
    headers = {"authorization": "Bearer sk-ant-oat01-abc123def456"}
    assert classify_auth_mode(headers) is AuthMode.OAUTH


def test_claude_code_ua_classified_subscription() -> None:
    """Claude Code CLI: ``User-Agent: claude-code/1.2.3 ...``."""
    headers = {"user-agent": "claude-code/1.2.3 (darwin; arm64)"}
    assert classify_auth_mode(headers) is AuthMode.SUBSCRIPTION


def test_cursor_ua_classified_subscription() -> None:
    """Cursor CLI: ``User-Agent: cursor/1.0``."""
    headers = {"user-agent": "cursor/1.0"}
    assert classify_auth_mode(headers) is AuthMode.SUBSCRIPTION


def test_no_auth_no_user_agent_default_payg() -> None:
    """Empty headers → safest default is PAYG.

    The OAuth/bedrock branch fires only when there's a positive
    non-Bearer auth signal. Choosing PAYG by default favors the
    OSS-default workload (per-token cost saving).
    """
    assert classify_auth_mode({}) is AuthMode.PAYG


def test_bedrock_no_auth_classified_oauth() -> None:
    """Bedrock SigV4: ``Authorization: AWS4-HMAC-SHA256 Credential=...``.

    Not a Bearer scheme; we treat all non-Bearer Authorization as
    OAuth (passthrough-prefer).
    """
    headers = {
        "authorization": (
            "AWS4-HMAC-SHA256 Credential=AKIAIOSFODNN7EXAMPLE/20260501/"
            "us-east-1/bedrock/aws4_request, SignedHeaders=host;x-amz-date, "
            "Signature=fe5f80f77d5fa3beca038a248ff027"
        ),
    }
    assert classify_auth_mode(headers) is AuthMode.OAUTH


# ── Bonus matrix ──────────────────────────────────────────────────


def test_openai_payg_sk_classified_payg() -> None:
    """OpenAI PAYG: ``Authorization: Bearer sk-proj-...``."""
    headers = {"authorization": "Bearer sk-proj-abcdef0123456789"}
    assert classify_auth_mode(headers) is AuthMode.PAYG


def test_gemini_x_goog_api_key_classified_payg() -> None:
    """Google Gemini API key as ``x-goog-api-key``."""
    headers = {"x-goog-api-key": "AIzaSyDUMMYKEY1234567890"}
    assert classify_auth_mode(headers) is AuthMode.PAYG


def test_subscription_takes_precedence_over_oauth_token() -> None:
    """Claude Code CLI sends ``Bearer sk-ant-oat-...`` but is subscription.

    UA wins because the rate-limit / fingerprint policy is bound to
    the CLI, not the token shape. The same token via a non-CLI UA
    would be OAuth.
    """
    headers = {
        "user-agent": "claude-code/1.5.0 (linux; x86_64)",
        "authorization": "Bearer sk-ant-oat-01-abc123",
    }
    assert classify_auth_mode(headers) is AuthMode.SUBSCRIPTION


# ── Edge cases (defensive coverage) ──────────────────────────────


def test_anthropic_x_api_key_classified_payg() -> None:
    """Anthropic API key style: ``x-api-key: sk-ant-...``."""
    headers = {"x-api-key": "sk-ant-api03-abcdef"}
    assert classify_auth_mode(headers) is AuthMode.PAYG


@pytest.mark.parametrize("prefix", SUBSCRIPTION_UA_PREFIXES)
def test_every_subscription_prefix_classified_subscription(prefix: str) -> None:
    """Each entry in the prefix list classifies as Subscription on its own."""
    headers = {"user-agent": f"{prefix}1.0 (test)"}
    assert classify_auth_mode(headers) is AuthMode.SUBSCRIPTION


def test_unparseable_authorization_does_not_raise() -> None:
    """Bytes-valued non-UTF-8 authorization warns and falls through.

    Mirrors the Rust ``classify`` behaviour: never panic, always
    return a valid ``AuthMode``. With nothing else to disambiguate,
    we land on PAYG (the safe default).
    """
    headers = {"authorization": b"\xffnope"}
    # Should not raise.
    result = classify_auth_mode(headers)
    assert result is AuthMode.PAYG


def test_case_insensitive_header_lookup() -> None:
    """Header names are matched case-insensitively (Starlette parity)."""
    headers = {"Authorization": "Bearer sk-ant-api03-test"}
    assert classify_auth_mode(headers) is AuthMode.PAYG


def test_enum_values_match_rust_as_str() -> None:
    """The string form of each enum member matches the Rust ``as_str()``."""
    assert AuthMode.PAYG.value == "payg"
    assert AuthMode.OAUTH.value == "oauth"
    assert AuthMode.SUBSCRIPTION.value == "subscription"


# ── Performance ──────────────────────────────────────────────────


def test_classify_under_100us_per_call() -> None:
    """Smoke perf check.

    Python is slower than Rust by an order of magnitude on string
    work, so the budget here is 100us (vs 10us in Rust). The proxy
    only calls this once per request, and the absolute ceiling is
    well under 1ms; this guards against pathological regressions.
    """
    headers = {
        "user-agent": "claude-code/1.5.0 (linux; x86_64) anthropic/0.42.0",
        "authorization": "Bearer sk-ant-oat-01-abcdefghijklmnopqrstuv",
        "content-type": "application/json",
        "accept": "application/json",
        "host": "api.anthropic.com",
    }

    # Warmup
    for _ in range(1000):
        classify_auth_mode(headers)

    iters = 10_000
    start = time.perf_counter()
    for _ in range(iters):
        classify_auth_mode(headers)
    elapsed = time.perf_counter() - start
    per_call_us = (elapsed / iters) * 1_000_000

    assert per_call_us < 100, f"classify_auth_mode took {per_call_us:.2f} us/call (limit: 100 us)"


def test_classify_client_uses_default_when_no_client_signal():
    headers = {"user-agent": "anthropic/0.42.0"}

    assert classify_client(headers, default="claude") == "claude"