atlas-iac/testing/tests/test_hermes_ai_usage_exporter.py

"""Tests for the credential-safe Atlas AI usage exporter."""

from __future__ import annotations

import importlib.util
import json
import sys
from datetime import date
from pathlib import Path


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


def load_module():
    """Load the standalone exporter script as a test module."""
    spec = importlib.util.spec_from_file_location("ai_usage_exporter", SCRIPT)
    module = importlib.util.module_from_spec(spec)
    assert spec and spec.loader
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


def test_codex_payloads_include_real_windows_and_bounded_token_periods():
    mod = load_module()
    rate = {
        "rateLimitsByLimitId": {
            "codex": {
                "primary": {
                    "usedPercent": 41,
                    "windowDurationMins": 10080,
                    "resetsAt": 1_800_000_000,
                }
            },
            "opaque_backend_id": {
                "limitName": "GPT-5.3-Codex-Spark",
                "primary": {
                    "usedPercent": 0,
                    "windowDurationMins": 10080,
                    "resetsAt": 1_800_100_000,
                },
            },
        }
    }
    usage = {
        "summary": {"lifetimeTokens": 1234, "peakDailyTokens": 500},
        "dailyUsageBuckets": [
            {"startDate": "2026-08-16", "tokens": 100},
            {"startDate": "2026-08-15", "tokens": 50},
            {"startDate": "2026-07-01", "tokens": 900},
        ],
    }

    samples = mod.parse_codex_payloads(rate, usage, today=date(2026, 8, 16))
    values = {(item.name, tuple(sorted(item.labels.items()))): item.value for item in samples}

    codex_labels = (
        ("limit", "codex"),
        ("provider", "openai"),
        ("window", "seven_day"),
    )
    spark_labels = (
        ("limit", "gpt-5-3-codex-spark"),
        ("provider", "openai"),
        ("window", "seven_day"),
    )
    assert values[("atlas_ai_quota_used_percent", codex_labels)] == 41
    assert values[("atlas_ai_quota_remaining_percent", codex_labels)] == 59
    assert values[("atlas_ai_quota_used_percent", spark_labels)] == 0
    assert values[
        (
            "atlas_ai_account_tokens",
            (("period", "latest_day"), ("provider", "openai")),
        )
    ] == 100
    assert values[
        (
            "atlas_ai_account_tokens",
            (("period", "seven_day"), ("provider", "openai")),
        )
    ] == 150


def test_claude_payload_uses_only_supported_quota_fields():
    mod = load_module()
    secret = "secret-access-token-must-not-leak"
    payload = {
        "five_hour": {
            "utilization": 16.0,
            "resets_at": "2026-08-16T08:39:59+00:00",
            "unexpected_secret": secret,
        },
        "seven_day": {
            "utilization": 26.0,
            "resets_at": "2026-08-20T17:59:59+00:00",
        },
        "internal_bucket": {"utilization": 99, "resets_at": secret},
        "extra_usage": {"is_enabled": False, "used_credits": secret},
    }

    samples = mod.parse_claude_payload(payload)
    assert {item.labels["window"] for item in samples if "window" in item.labels} == {
        "five_hour",
        "seven_day",
    }
    assert any(
        item.name == "atlas_ai_quota_remaining_percent"
        and item.labels["window"] == "seven_day"
        and item.value == 74
        for item in samples
    )
    assert secret not in repr(samples)


def test_render_reports_failure_and_freshness_without_logging_credentials():
    mod = load_module()
    collector = mod.Collector()
    state = collector._providers["openai"]
    state.samples = [
        mod.Sample(
            "atlas_ai_quota_used_percent",
            {"provider": "openai", "limit": "codex", "window": "seven_day"},
            42,
        )
    ]
    state.fetch_success = False
    state.last_attempt = 200
    state.last_success = 100
    rendered = collector.render().decode()

    assert 'atlas_ai_quota_fetch_success{provider="openai"} 0' in rendered
    assert 'atlas_ai_quota_last_success_timestamp_seconds{provider="openai"} 100' in rendered
    assert 'atlas_ai_quota_used_percent{limit="codex",provider="openai",window="seven_day"} 42' in rendered
    assert "accessToken" not in rendered
    assert "refreshToken" not in rendered


def test_codex_query_uses_structured_app_server_protocol(tmp_path, monkeypatch):
    mod = load_module()
    request_log = tmp_path / "requests.jsonl"
    mock = tmp_path / "codex-mock"
    mock.write_text(
        """#!/usr/bin/env python3
import json, os, sys
messages = [json.loads(sys.stdin.readline()) for _ in range(4)]
with open(os.environ['REQUEST_LOG'], 'w') as handle:
    for message in messages:
        handle.write(json.dumps(message) + '\\n')
responses = {
    1: {},
    2: {'rateLimits': {'primary': {'usedPercent': 10}}},
    3: {'summary': {}, 'dailyUsageBuckets': []},
}
for response_id, result in responses.items():
    print(json.dumps({'id': response_id, 'result': result}), flush=True)
"""
    )
    mock.chmod(0o755)
    monkeypatch.setattr(mod, "CODEX_BIN", str(mock))
    monkeypatch.setenv("REQUEST_LOG", str(request_log))

    rate, usage = mod.query_codex(timeout=5)
    requests = [json.loads(line) for line in request_log.read_text().splitlines()]

    assert rate["rateLimits"]["primary"]["usedPercent"] == 10
    assert usage["summary"] == {}
    assert [request["method"] for request in requests] == [
        "initialize",
        "initialized",
        "account/rateLimits/read",
        "account/usage/read",
    ]
    assert all("/status" not in json.dumps(request) for request in requests)