"""Tests for the credential-safe Atlas AI usage exporter.""" from __future__ import annotations import importlib.util import sys from datetime import date from io import BytesIO from pathlib import Path import yaml ROOT = Path(__file__).resolve().parents[2] SCRIPT = ROOT / "services/hermes/scripts/ai_usage_exporter.py" HERMES = ROOT / "services/hermes" def load_module(): """Load the standalone exporter script as a test module.""" script_directory = str(SCRIPT.parent) if script_directory not in sys.path: sys.path.insert(0, script_directory) 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", }, "seven_day_fable": { "utilization": 81.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", "seven_day_fable", } assert any( item.name == "atlas_ai_quota_remaining_percent" and item.labels["window"] == "seven_day" and item.value == 74 for item in samples ) assert any( item.name == "atlas_ai_quota_remaining_percent" and item.labels["window"] == "seven_day_fable" and item.value == 19 for item in samples ) assert secret not in repr(samples) def test_render_reports_failure_and_freshness_without_logging_credentials( monkeypatch, ): mod = load_module() monkeypatch.setattr( mod.claude_query, "credential_refresh_expiry_timestamp", lambda: 1_800_000_000, ) 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 ( 'atlas_ai_quota_credential_refresh_expiry_timestamp_seconds{provider="anthropic"} 1800000000' in rendered ) assert "accessToken" not in rendered assert "refreshToken" not in rendered def test_unexpected_provider_failure_is_sanitized_and_preserves_last_good( monkeypatch, capsys ): mod = load_module() collector = mod.Collector() previous = mod.Sample( "atlas_ai_quota_used_percent", {"provider": "openai", "limit": "codex", "window": "seven_day"}, 42, ) collector._providers["openai"].samples = [previous] secret = "provider-body-must-not-appear" class UnexpectedProviderFailure(Exception): pass def fail(): raise UnexpectedProviderFailure(secret) monkeypatch.setattr(mod, "query_codex", fail) collector.refresh_provider("openai") output = capsys.readouterr().out state = collector._providers["openai"] assert state.samples == [previous] assert state.fetch_success is False assert state.last_attempt > 0 assert "UnexpectedProviderFailure" in output assert secret not in output def test_polling_engine_isolates_providers_and_future_cycles(monkeypatch, capsys): mod = load_module() collector = mod.Collector() calls = [] secret = "cleanup-detail-must-not-appear" class UnexpectedCleanupFailure(Exception): pass def refresh(provider): calls.append(provider) if provider == "openai" and calls.count("openai") == 1: raise UnexpectedCleanupFailure(secret) monkeypatch.setattr(collector, "refresh_provider", refresh) engine = mod.PollingEngine(collector, interval=60) engine.poll_once() engine.poll_once() output = capsys.readouterr().out assert calls == ["openai", "anthropic", "openai", "anthropic"] assert collector._providers["openai"].fetch_success is False assert collector._providers["openai"].last_attempt > 0 assert "UnexpectedCleanupFailure" in output assert secret not in output def test_health_tracks_poller_liveness_and_progress_not_provider_success(): mod = load_module() now = [100.0] collector = mod.Collector() collector._providers["openai"].fetch_success = False default_engine = mod.PollingEngine(collector, interval=300, clock=lambda: now[0]) assert default_engine.progress_timeout == 420 engine = mod.PollingEngine( collector, interval=60, startup_grace=10, progress_timeout=75, clock=lambda: now[0], ) class ThreadState: alive = True def is_alive(self): return self.alive thread = ThreadState() engine._thread = thread assert engine.is_healthy() now[0] = 111.0 assert not engine.is_healthy() engine._mark_progress() assert engine.is_healthy() now[0] = 187.0 assert not engine.is_healthy() thread.alive = False now[0] = 100.0 assert not engine.is_healthy() def test_manifest_rolls_out_the_bounded_codex_deadline_and_poller_module(): """Flux must restart the exporter with every reliability input mounted.""" deployment = yaml.safe_load((HERMES / "agent-deployment.yaml").read_text()) annotations = deployment["spec"]["template"]["metadata"]["annotations"] containers = deployment["spec"]["template"]["spec"]["containers"] exporter = next(item for item in containers if item["name"] == "ai-usage-exporter") environment = {item["name"]: item["value"] for item in exporter["env"]} assert annotations["ai.bstein.dev/config-rev"] == ( "20260825-claude-quota-expiry" ) assert environment["ATLAS_AI_CODEX_QUERY_TIMEOUT_SECONDS"] == "45" assert environment["ATLAS_AI_AUTHENTICATION_GRACE_SECONDS"] == "1200" assert environment["ATLAS_AI_CLAUDE_OAUTH_TOKEN_FILE"] == ( "/claude-oauth-access/token" ) assert environment["ATLAS_AI_CLAUDE_CREDENTIALS"] == ( "/runtime-access/claude/.credentials.json" ) assert environment["ATLAS_AI_CLAUDE_USAGE_URL"] == ( "https://api.anthropic.com/api/oauth/usage" ) assert environment["ATLAS_AI_CLAUDE_QUOTA_MODEL"] == ( "claude-haiku-4-5-20251001" ) mounts = {item["name"]: item for item in exporter["volumeMounts"]} assert mounts["claude-oauth-access"]["readOnly"] is True kustomization = (HERMES / "kustomization.yaml").read_text() assert "ai_usage_claude.py=scripts/ai_usage_claude.py" in kustomization assert "ai_usage_polling.py=scripts/ai_usage_polling.py" in kustomization def test_health_endpoint_uses_poller_state_and_ignores_provider_failure(): mod = load_module() collector = mod.Collector() class Poller: healthy = False def is_healthy(self): return self.healthy poller = Poller() handler_class = mod.make_handler(collector, poller) handler = object.__new__(handler_class) handler.path = "/healthz" statuses = [] handler.send_response = statuses.append handler.send_header = lambda _name, _value: None handler.end_headers = lambda: None handler.wfile = BytesIO() handler.do_GET() assert statuses == [503] poller.healthy = True statuses.clear() collector._providers["openai"].fetch_success = False handler.do_GET() assert statuses == [200] def test_metrics_response_tolerates_disconnected_client(): mod = load_module() collector = mod.Collector() class HealthyPoller: def is_healthy(self): return True class Disconnected(BytesIO): def write(self, _body): raise BrokenPipeError handler_class = mod.make_handler(collector, HealthyPoller()) handler = object.__new__(handler_class) handler.path = "/metrics" handler.send_response = lambda _status: None handler.send_header = lambda _name, _value: None handler.end_headers = lambda: None handler.wfile = Disconnected() handler.do_GET()