394 lines
12 KiB
Python
394 lines
12 KiB
Python
"""Tests for the credential-safe Atlas AI usage exporter."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import json
|
|
import subprocess
|
|
import sys
|
|
from datetime import date
|
|
from io import BytesIO, StringIO
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
|
|
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)
|
|
|
|
|
|
def test_codex_cleanup_is_bounded_when_process_cannot_be_reaped(monkeypatch):
|
|
mod = load_module()
|
|
processes = []
|
|
selectors = []
|
|
|
|
class StubbornProcess:
|
|
def __init__(self, *_args, **_kwargs):
|
|
self.stdin = StringIO()
|
|
self.stdout = StringIO(
|
|
"\n".join(
|
|
json.dumps({"id": response_id, "result": result})
|
|
for response_id, result in (
|
|
(1, {}),
|
|
(2, {"rateLimits": {}}),
|
|
(3, {"summary": {}, "dailyUsageBuckets": []}),
|
|
)
|
|
)
|
|
+ "\n"
|
|
)
|
|
self.terminate_calls = 0
|
|
self.kill_calls = 0
|
|
self.wait_timeouts = []
|
|
processes.append(self)
|
|
|
|
def terminate(self):
|
|
self.terminate_calls += 1
|
|
|
|
def kill(self):
|
|
self.kill_calls += 1
|
|
|
|
def wait(self, timeout):
|
|
self.wait_timeouts.append(timeout)
|
|
raise subprocess.TimeoutExpired("codex-mock", timeout)
|
|
|
|
class TrackingSelector:
|
|
def __init__(self):
|
|
self.fileobj = None
|
|
self.closed = False
|
|
selectors.append(self)
|
|
|
|
def register(self, fileobj, _events):
|
|
self.fileobj = fileobj
|
|
|
|
def select(self, _timeout):
|
|
return [(SimpleNamespace(fileobj=self.fileobj), None)]
|
|
|
|
def close(self):
|
|
self.closed = True
|
|
|
|
monkeypatch.setattr(mod.subprocess, "Popen", StubbornProcess)
|
|
monkeypatch.setattr(mod.selectors, "DefaultSelector", TrackingSelector)
|
|
|
|
try:
|
|
mod.query_codex(timeout=1)
|
|
except Exception as error: # The cleanup result must be sanitized and non-TimeoutExpired.
|
|
assert type(error) is RuntimeError
|
|
assert str(error) == "Codex app-server cleanup failed"
|
|
else:
|
|
raise AssertionError("unreaped Codex process was reported as successful")
|
|
|
|
process = processes[0]
|
|
assert process.terminate_calls == 1
|
|
assert process.kill_calls == 1
|
|
assert process.wait_timeouts == [5, 5]
|
|
assert process.stdin.closed
|
|
assert process.stdout.closed
|
|
assert selectors[0].closed
|
|
|
|
|
|
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
|
|
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_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()
|