atlas-iac/testing/tests/test_hermes_ai_usage_claude_scoped.py

"""Scoped Claude Code quota collection and credential-safety contracts."""

from __future__ import annotations

import importlib.util
import json
import stat
import sys
from pathlib import Path
from urllib.error import HTTPError

import pytest


ROOT = Path(__file__).resolve().parents[2]
SCRIPT = ROOT / "services/hermes/scripts/ai_usage_claude.py"


def load_module():
    """Load a fresh Claude usage helper for environment-independent tests."""
    name = "ai_usage_claude_scoped"
    spec = importlib.util.spec_from_file_location(name, SCRIPT)
    module = importlib.util.module_from_spec(spec)
    assert spec and spec.loader
    sys.modules[name] = module
    spec.loader.exec_module(module)
    return module


class Response:
    """Small bounded urllib response fixture."""

    def __init__(self, body, *, headers=None):
        self.body = body if isinstance(body, bytes) else json.dumps(body).encode()
        self.headers = headers or {}

    def __enter__(self):
        return self

    def __exit__(self, *_args):
        return None

    def read(self, maximum):
        assert maximum == 1 << 20
        return self.body


def write_credentials(path, **overrides):
    """Write a private synthetic Claude OAuth document."""
    credentials = {
        "accessToken": "scoped-access",
        "refreshToken": "scoped-refresh",
        "expiresAt": 2_000_000,
        "scopes": [
            "user:profile",
            "user:inference",
            "user:sessions:claude_code",
        ],
    }
    credentials.update(overrides)
    path.write_text(
        json.dumps({"claudeAiOauth": credentials}) + "\n",
        encoding="utf-8",
    )
    path.chmod(0o600)


def test_scoped_payload_selects_only_overall_and_exact_fable_limits():
    mod = load_module()
    secret = "provider-document-detail-must-not-leak"

    payload = mod._scoped_payload(
        {
            "five_hour": {"utilization": 58, "resets_at": "2026-08-24T17:20:00Z"},
            "seven_day": {"utilization": 59, "resets_at": "2026-08-27T18:00:00Z"},
            "limits": [
                {
                    "kind": "weekly_scoped",
                    "percent": 81,
                    "resets_at": "2026-08-27T18:00:00Z",
                    "scope": {"model": {"display_name": "Fable"}},
                    "private": secret,
                },
                {
                    "kind": "weekly_scoped",
                    "percent": 90,
                    "scope": {"model": {"display_name": "Unknown"}},
                },
            ],
            "account": secret,
        }
    )

    assert payload == {
        "five_hour": {
            "utilization": 58.0,
            "resets_at": "2026-08-24T17:20:00Z",
        },
        "seven_day": {
            "utilization": 59.0,
            "resets_at": "2026-08-27T18:00:00Z",
        },
        "seven_day_fable": {
            "utilization": 81.0,
            "resets_at": "2026-08-27T18:00:00Z",
        },
    }
    assert secret not in repr(payload)


def test_scoped_query_uses_current_private_access_token(tmp_path, monkeypatch):
    mod = load_module()
    credentials = tmp_path / ".credentials.json"
    write_credentials(credentials)
    monkeypatch.setattr(mod, "CREDENTIALS_FILE", credentials)
    monkeypatch.setattr(mod.time, "time", lambda: 1_000)
    requests = []

    def open_request(request, timeout):
        requests.append(request)
        assert timeout == 30
        return Response(
            {
                "seven_day": {"utilization": 42},
                "limits": [
                    {
                        "kind": "weekly_scoped",
                        "percent": 73,
                        "scope": {"model": {"display_name": "Fable 5"}},
                    }
                ],
            }
        )

    monkeypatch.setattr(mod, "urlopen", open_request)

    assert mod._query_scoped_usage() == {
        "seven_day": {"utilization": 42.0},
        "seven_day_fable": {"utilization": 73.0},
    }
    assert len(requests) == 1
    assert requests[0].full_url == mod.USAGE_URL
    assert requests[0].data is None
    assert requests[0].get_header("Authorization") == "Bearer scoped-access"


def test_credential_refresh_expiry_exposes_only_a_timestamp(tmp_path, monkeypatch):
    mod = load_module()
    credentials = tmp_path / ".credentials.json"
    write_credentials(credentials, refreshTokenExpiresAt=2_000_000)
    monkeypatch.setattr(mod, "CREDENTIALS_FILE", credentials)

    assert mod.credential_refresh_expiry_timestamp() == 2_000

    write_credentials(credentials, refreshTokenExpiresAt=True)
    assert mod.credential_refresh_expiry_timestamp() is None
    credentials.unlink()
    assert mod.credential_refresh_expiry_timestamp() is None


def test_expired_scoped_token_refreshes_and_persists_rotation(
    tmp_path, monkeypatch
):
    mod = load_module()
    credentials = tmp_path / ".credentials.json"
    write_credentials(credentials, expiresAt=900_000)
    monkeypatch.setattr(mod, "CREDENTIALS_FILE", credentials)
    monkeypatch.setattr(mod.time, "time", lambda: 1_000)
    requests = []

    def open_request(request, timeout):
        requests.append(request)
        assert timeout == 30
        if request.full_url == mod.TOKEN_URL:
            return Response(
                {
                    "access_token": "rotated-access",
                    "refresh_token": "rotated-refresh",
                    "expires_in": 3600,
                    "refresh_token_expires_in": 7200,
                    "scope": "user:profile user:inference",
                }
            )
        return Response(
            {
                "seven_day": {"utilization": 50},
                "limits": [
                    {
                        "kind": "weekly_scoped",
                        "percent": 75,
                        "scope": {"model": {"display_name": "Fable"}},
                    }
                ],
            }
        )

    monkeypatch.setattr(mod, "urlopen", open_request)

    assert mod._query_scoped_usage()["seven_day_fable"]["utilization"] == 75
    refresh_request = requests[0]
    assert refresh_request.full_url == mod.TOKEN_URL
    refresh_body = json.loads(refresh_request.data)
    assert refresh_body == {
        "grant_type": "refresh_token",
        "refresh_token": "scoped-refresh",
        "client_id": mod.OAUTH_CLIENT_ID,
        "scope": "user:profile user:inference user:sessions:claude_code",
    }
    persisted = json.loads(credentials.read_text())["claudeAiOauth"]
    assert persisted["accessToken"] == "rotated-access"
    assert persisted["refreshToken"] == "rotated-refresh"
    assert persisted["expiresAt"] == 4_600_000
    assert persisted["refreshTokenExpiresAt"] == 8_200_000
    assert stat.S_IMODE(credentials.stat().st_mode) == 0o600


def test_query_falls_back_to_setup_token_headers_when_scope_is_unavailable(
    tmp_path, monkeypatch
):
    mod = load_module()
    missing_credentials = tmp_path / "missing-credentials"
    setup_token = tmp_path / "setup-token"
    setup_token.write_text("setup-only-token\n", encoding="utf-8")
    monkeypatch.setattr(mod, "CREDENTIALS_FILE", missing_credentials)
    monkeypatch.setattr(mod, "OAUTH_TOKEN_FILE", setup_token)
    requests = []

    def open_request(request, timeout):
        requests.append(request)
        assert timeout == 30
        return Response(
            {},
            headers={
                "anthropic-ratelimit-unified-5h-utilization": "0.25",
                "anthropic-ratelimit-unified-7d-utilization": "0.40",
            },
        )

    monkeypatch.setattr(mod, "urlopen", open_request)

    assert mod.query_claude() == {
        "five_hour": {"utilization": 25.0},
        "seven_day": {"utilization": 40.0},
    }
    assert len(requests) == 1
    assert requests[0].full_url == mod.MESSAGES_URL
    assert requests[0].get_header("Authorization") == "Bearer setup-only-token"


def test_scoped_validation_and_unauthorized_retry_fail_closed(
    tmp_path, monkeypatch
):
    mod = load_module()
    credentials = tmp_path / ".credentials.json"
    write_credentials(credentials)
    monkeypatch.setattr(mod, "CREDENTIALS_FILE", credentials)
    monkeypatch.setattr(mod.time, "time", lambda: 1_000)

    assert mod._used_percentage(True) is None
    assert mod._used_percentage(-1) is None
    assert mod._used_percentage(101) is None
    with pytest.raises(RuntimeError, match="supported quota windows"):
        mod._scoped_payload({"limits": [{"kind": "weekly_scoped"}]})

    attempts = []

    def open_request(request, timeout):
        attempts.append(request.full_url)
        if len(attempts) == 1:
            raise HTTPError(request.full_url, 401, "secret-body", {}, None)
        if request.full_url == mod.TOKEN_URL:
            return Response(
                {
                    "access_token": "retried-access",
                    "refresh_token": "retried-refresh",
                    "expires_in": 3600,
                }
            )
        return Response({"seven_day": {"utilization": 12}})

    monkeypatch.setattr(mod, "urlopen", open_request)
    assert mod._query_scoped_usage() == {"seven_day": {"utilization": 12.0}}
    assert attempts == [mod.USAGE_URL, mod.TOKEN_URL, mod.USAGE_URL]


def test_scoped_payload_ignores_malformed_and_unrelated_limit_entries():
    mod = load_module()

    assert mod._scoped_payload(
        {
            "five_hour": {"utilization": "invalid"},
            "seven_day": {"utilization": 10},
            "limits": [
                None,
                {"kind": "session", "percent": 20},
                {
                    "kind": "weekly_scoped",
                    "percent": 30,
                    "scope": {"model": {"display_name": "Sonnet"}},
                },
                {
                    "kind": "weekly_scoped",
                    "percent": "invalid",
                    "scope": {"model": {"display_name": "Fable"}},
                },
            ],
        }
    ) == {"seven_day": {"utilization": 10.0}}
    with pytest.raises(RuntimeError, match="not an object"):
        mod._scoped_payload(None)


@pytest.mark.parametrize(
    ("content", "mode", "message"),
    [
        ('{"claudeAiOauth":{}}\n', 0o644, "unsafe"),
        ("not-json\n", 0o600, "invalid"),
        ("[]\n", 0o600, "invalid"),
    ],
)
def test_private_credential_reader_rejects_unsafe_documents(
    tmp_path, monkeypatch, content, mode, message
):
    mod = load_module()
    credentials = tmp_path / ".credentials.json"
    credentials.write_text(content, encoding="utf-8")
    credentials.chmod(mode)
    monkeypatch.setattr(mod, "CREDENTIALS_FILE", credentials)

    with pytest.raises(RuntimeError, match=message):
        mod._read_credentials()


def test_credential_writer_and_refresh_validation_are_bounded(tmp_path, monkeypatch):
    mod = load_module()
    credentials = tmp_path / ".credentials.json"
    write_credentials(credentials)
    monkeypatch.setattr(mod, "CREDENTIALS_FILE", credentials)
    monkeypatch.setattr(mod, "MAX_DOCUMENT_BYTES", 1)
    with pytest.raises(RuntimeError, match="too large"):
        mod._write_credentials({"claudeAiOauth": {"accessToken": "large"}})

    monkeypatch.setattr(mod, "MAX_DOCUMENT_BYTES", 1 << 20)
    with pytest.raises(RuntimeError, match="refresh credential is unavailable"):
        mod._refresh_access_token({"claudeAiOauth": {}}, {})

    document = {
        "claudeAiOauth": {
            "accessToken": "old-access",
            "refreshToken": "old-refresh",
            "scopes": ["unapproved:scope"],
        }
    }
    for response_document in ([], {"access_token": "", "expires_in": True}):
        monkeypatch.setattr(
            mod,
            "urlopen",
            lambda request, timeout, value=response_document: Response(value),
        )
        with pytest.raises(RuntimeError, match="refresh response is invalid"):
            mod._refresh_access_token(document, document["claudeAiOauth"])


def test_non_authentication_usage_error_is_not_retried(tmp_path, monkeypatch):
    mod = load_module()
    credentials = tmp_path / ".credentials.json"
    write_credentials(credentials)
    monkeypatch.setattr(mod, "CREDENTIALS_FILE", credentials)
    monkeypatch.setattr(mod.time, "time", lambda: 1_000)
    attempts = []

    def reject(request, timeout):
        attempts.append(request.full_url)
        raise HTTPError(request.full_url, 503, "provider-detail", {}, None)

    monkeypatch.setattr(mod, "urlopen", reject)
    with pytest.raises(HTTPError):
        mod._query_scoped_usage()
    assert attempts == [mod.USAGE_URL]