"""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)