fe7832efaf
Some checks failed
Tests / Declarative: Post Actions failed: 39, skipped: 19, passed: 2757
334 lines
10 KiB
Python
334 lines
10 KiB
Python
"""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",
|
|
},
|
|
"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_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"] == (
|
|
"20260823-dual-provider-quota-health"
|
|
)
|
|
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_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()
|