376 lines
12 KiB
Python
376 lines
12 KiB
Python
"""Focused edge coverage for AI usage exporter provider and polling boundaries."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import json
|
|
import sys
|
|
from io import BytesIO
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[2]
|
|
SCRIPT = ROOT / "services/hermes/scripts/ai_usage_exporter.py"
|
|
|
|
|
|
def load_module():
|
|
"""Load a fresh exporter while making its sibling modules importable."""
|
|
script_directory = str(SCRIPT.parent)
|
|
if script_directory not in sys.path:
|
|
sys.path.insert(0, script_directory)
|
|
name = "ai_usage_exporter_coverage"
|
|
spec = importlib.util.spec_from_file_location(name, SCRIPT)
|
|
module = importlib.util.module_from_spec(spec)
|
|
assert spec and spec.loader
|
|
sys.modules[name] = module
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
def test_parser_rejects_nonfinite_and_malformed_optional_fields():
|
|
mod = load_module()
|
|
|
|
assert mod._number(True) is None
|
|
assert mod._number("1") is None
|
|
assert mod._number(float("nan")) is None
|
|
assert mod._number(float("inf")) is None
|
|
assert mod._timestamp(None) is None
|
|
assert mod._timestamp("not-a-time") is None
|
|
assert mod._window_name({}, "fallback") == "fallback"
|
|
assert mod._codex_limit_name("other", {}) == "additional"
|
|
|
|
samples = mod.parse_codex_payloads(
|
|
{
|
|
"rateLimitsByLimitId": {
|
|
1: {},
|
|
"not-a-snapshot": "invalid",
|
|
"unnamed": {
|
|
"primary": {"usedPercent": None},
|
|
"secondary": {
|
|
"usedPercent": 120,
|
|
"resetsAt": "invalid",
|
|
},
|
|
},
|
|
}
|
|
},
|
|
{
|
|
"summary": {
|
|
"lifetimeTokens": None,
|
|
"currentStreakDays": 2,
|
|
"longestStreakDays": 3,
|
|
"longestRunningTurnSec": 4,
|
|
},
|
|
"dailyUsageBuckets": [
|
|
None,
|
|
{"startDate": "invalid", "tokens": 2},
|
|
{"startDate": "2026-08-17", "tokens": None},
|
|
],
|
|
},
|
|
)
|
|
|
|
assert any(sample.value == 120 for sample in samples)
|
|
assert not any(
|
|
sample.name == "atlas_ai_quota_reset_timestamp_seconds" for sample in samples
|
|
)
|
|
assert mod.parse_codex_payloads({"rateLimits": None}, {"summary": None})
|
|
|
|
|
|
def test_claude_parser_ignores_unbounded_fields_and_invalid_extra_usage():
|
|
mod = load_module()
|
|
|
|
assert mod.parse_claude_payload(
|
|
{
|
|
"five_hour": {"utilization": None},
|
|
"seven_day": {"utilization": 10, "resets_at": None},
|
|
"extra_usage": {"is_enabled": "yes"},
|
|
}
|
|
)
|
|
assert mod.parse_claude_payload({"extra_usage": "invalid"}) == []
|
|
|
|
|
|
def test_claude_header_probe_rejects_malformed_optional_values(
|
|
tmp_path, monkeypatch
|
|
):
|
|
mod = load_module()
|
|
missing = tmp_path / "missing-token"
|
|
monkeypatch.setattr(mod.claude_query, "OAUTH_TOKEN_FILE", missing)
|
|
|
|
with pytest.raises(RuntimeError, match="token is unavailable"):
|
|
mod.claude_query._setup_token()
|
|
|
|
assert mod.claude_query._percentage("invalid") is None
|
|
assert mod.claude_query._percentage(None) is None
|
|
assert mod.claude_query._percentage("-0.1") is None
|
|
assert mod.claude_query._percentage("1.1") is None
|
|
assert mod.claude_query._timestamp("invalid") is None
|
|
assert mod.claude_query._timestamp(None) is None
|
|
assert mod.claude_query._timestamp("0") is None
|
|
assert mod.claude_query._header_payload(
|
|
{"anthropic-ratelimit-unified-5h-utilization": "0.5"}
|
|
) == {"five_hour": {"utilization": 50}}
|
|
|
|
|
|
def test_provider_queries_validate_credentials_and_response_shape(
|
|
tmp_path, monkeypatch
|
|
):
|
|
mod = load_module()
|
|
token_file = tmp_path / "claude-oauth-token"
|
|
token_file.write_text("", encoding="utf-8")
|
|
monkeypatch.setattr(mod.claude_query, "OAUTH_TOKEN_FILE", token_file)
|
|
|
|
with pytest.raises(RuntimeError, match="token is unavailable"):
|
|
mod.query_claude()
|
|
|
|
token_file.write_text("runtime-only-token\n", encoding="utf-8")
|
|
response_headers = iter(
|
|
[
|
|
{
|
|
"anthropic-ratelimit-unified-5h-utilization": "0.16",
|
|
"anthropic-ratelimit-unified-5h-reset": "1800000000",
|
|
"anthropic-ratelimit-unified-7d-utilization": "0.26",
|
|
"anthropic-ratelimit-unified-7d-reset": "1800100000",
|
|
},
|
|
{},
|
|
]
|
|
)
|
|
requests = []
|
|
|
|
class Response:
|
|
def __init__(self):
|
|
self.headers = next(response_headers)
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *_args):
|
|
return None
|
|
|
|
def read(self, maximum):
|
|
assert maximum == 1 << 20
|
|
return b"{}"
|
|
|
|
def open_request(request, timeout):
|
|
requests.append(request)
|
|
assert timeout == 30
|
|
return Response()
|
|
|
|
monkeypatch.setattr(mod.claude_query, "urlopen", open_request)
|
|
assert mod.query_claude() == {
|
|
"five_hour": {"utilization": 16, "resets_at": 1_800_000_000},
|
|
"seven_day": {"utilization": 26, "resets_at": 1_800_100_000},
|
|
}
|
|
assert requests[0].get_header("Authorization") == "Bearer runtime-only-token"
|
|
assert requests[0].get_header("Anthropic-beta") == (
|
|
"oauth-2025-04-20,claude-code-20250219"
|
|
)
|
|
assert json.loads(requests[0].data) == {
|
|
"model": "claude-haiku-4-5-20251001",
|
|
"max_tokens": 1,
|
|
"messages": [{"role": "user", "content": "quota"}],
|
|
}
|
|
with pytest.raises(RuntimeError, match="omitted subscription quota headers"):
|
|
mod.query_claude()
|
|
|
|
assert mod.query_codex is mod.codex_query.query_codex
|
|
|
|
|
|
def test_claude_quota_failure_does_not_log_the_vault_token(
|
|
tmp_path, monkeypatch, capsys
|
|
):
|
|
mod = load_module()
|
|
token_file = tmp_path / "claude-oauth-token"
|
|
token_file.write_text("revoked-runtime-token", encoding="utf-8")
|
|
monkeypatch.setattr(mod.claude_query, "OAUTH_TOKEN_FILE", token_file)
|
|
monkeypatch.setattr(
|
|
mod.claude_query,
|
|
"urlopen",
|
|
lambda *_args, **_kwargs: (_ for _ in ()).throw(
|
|
TimeoutError("revoked-runtime-token")
|
|
),
|
|
)
|
|
|
|
collector = mod.Collector()
|
|
collector.refresh_provider("anthropic")
|
|
|
|
assert collector._providers["anthropic"].fetch_success is False
|
|
assert "revoked-runtime-token" not in capsys.readouterr().out
|
|
|
|
|
|
def test_provider_authentication_uses_only_fresh_non_secret_health(tmp_path, monkeypatch):
|
|
mod = load_module()
|
|
monkeypatch.setattr(mod, "PROVIDER_HEALTH_ROOT", tmp_path)
|
|
health = tmp_path / "claude.json"
|
|
health.write_text('{"authenticated":true}\n', encoding="utf-8")
|
|
|
|
assert mod._provider_authenticated("anthropic")
|
|
monkeypatch.setattr(mod.time, "time", lambda: health.stat().st_mtime + 601)
|
|
assert not mod._provider_authenticated("anthropic")
|
|
|
|
|
|
def test_recent_quota_success_keeps_access_healthy_when_broker_snapshot_is_quiet(
|
|
monkeypatch,
|
|
):
|
|
mod = load_module()
|
|
collector = mod.Collector()
|
|
collector._providers["openai"].last_success = 1_000
|
|
monkeypatch.setattr(mod, "_provider_authenticated", lambda _provider: False)
|
|
monkeypatch.setattr(mod.time, "time", lambda: 1_300)
|
|
|
|
rendered = collector.render().decode()
|
|
|
|
assert 'atlas_ai_provider_authenticated{provider="openai"} 1' in rendered
|
|
|
|
monkeypatch.setattr(
|
|
mod.time,
|
|
"time",
|
|
lambda: 1_000 + mod.AUTHENTICATION_GRACE_SECONDS + 1,
|
|
)
|
|
rendered = collector.render().decode()
|
|
assert 'atlas_ai_provider_authenticated{provider="openai"} 0' in rendered
|
|
|
|
|
|
def test_collector_refreshes_both_providers_and_contains_unknown_provider(
|
|
monkeypatch, capsys
|
|
):
|
|
mod = load_module()
|
|
collector = mod.Collector()
|
|
monkeypatch.setattr(
|
|
mod,
|
|
"query_codex",
|
|
lambda: ({"rateLimits": {"primary": {"usedPercent": 1}}}, {}),
|
|
)
|
|
monkeypatch.setattr(mod, "query_claude", lambda: {"five_hour": {"utilization": 2}})
|
|
|
|
collector.refresh_provider("openai")
|
|
collector.refresh_provider("anthropic")
|
|
assert collector._providers["openai"].fetch_success
|
|
assert collector._providers["anthropic"].fetch_success
|
|
assert collector._providers["openai"].authenticated
|
|
assert collector._providers["anthropic"].authenticated
|
|
assert collector._providers["openai"].last_success > 0
|
|
|
|
collector.record_failure("openai")
|
|
assert not collector._providers["openai"].fetch_success
|
|
with pytest.raises(KeyError, match="unknown"):
|
|
collector.refresh_provider("unknown")
|
|
assert "unknown quota collection deferred: ValueError" in capsys.readouterr().out
|
|
|
|
|
|
def test_polling_failure_accounting_and_cycle_errors_are_sanitized(monkeypatch, capsys):
|
|
mod = load_module()
|
|
secret = "must-not-leak"
|
|
|
|
class Failure(Exception):
|
|
pass
|
|
|
|
class Collector:
|
|
def refresh_provider(self, _provider):
|
|
raise Failure(secret)
|
|
|
|
def record_failure(self, *_args, **_kwargs):
|
|
raise Failure(secret)
|
|
|
|
engine = mod.PollingEngine(Collector(), interval=60)
|
|
engine.poll_once()
|
|
output = capsys.readouterr().out
|
|
assert "quota failure accounting deferred: Failure" in output
|
|
assert secret not in output
|
|
|
|
monkeypatch.setattr(
|
|
engine, "poll_once", lambda: (_ for _ in ()).throw(Failure(secret))
|
|
)
|
|
monkeypatch.setattr(
|
|
mod.time, "sleep", lambda _interval: (_ for _ in ()).throw(StopIteration)
|
|
)
|
|
with pytest.raises(StopIteration):
|
|
engine.run()
|
|
output = capsys.readouterr().out
|
|
assert "quota polling cycle deferred: Failure" in output
|
|
assert secret not in output
|
|
|
|
|
|
def test_poller_start_is_idempotent_and_main_wires_server(monkeypatch):
|
|
mod = load_module()
|
|
threads = []
|
|
|
|
class Thread:
|
|
def __init__(self, **kwargs):
|
|
self.kwargs = kwargs
|
|
self.started = False
|
|
threads.append(self)
|
|
|
|
def start(self):
|
|
self.started = True
|
|
|
|
monkeypatch.setattr(mod.threading, "Thread", Thread)
|
|
engine = mod.PollingEngine(mod.Collector(), interval=60)
|
|
engine.start()
|
|
engine.start()
|
|
assert len(threads) == 1
|
|
assert threads[0].started
|
|
assert threads[0].kwargs["daemon"] is True
|
|
|
|
events = []
|
|
|
|
class Poller:
|
|
def __init__(self, collector, *, interval):
|
|
events.append(("poller", collector, interval))
|
|
|
|
def start(self):
|
|
events.append("started")
|
|
|
|
class Server:
|
|
def __init__(self, address, handler):
|
|
events.append(("server", address, handler))
|
|
|
|
def serve_forever(self):
|
|
events.append("served")
|
|
|
|
monkeypatch.setattr(mod, "PollingEngine", Poller)
|
|
monkeypatch.setattr(mod, "Server", Server)
|
|
monkeypatch.setattr(
|
|
mod, "make_handler", lambda collector, poller: (collector, poller)
|
|
)
|
|
monkeypatch.setenv("ATLAS_AI_USAGE_INTERVAL_SECONDS", "1")
|
|
monkeypatch.setenv("ATLAS_AI_USAGE_PORT", "19010")
|
|
|
|
assert mod.main() == 0
|
|
assert events[0][0] == "poller"
|
|
assert events[0][2] == 60
|
|
assert events[2][1][1] == 19010
|
|
assert events[-1] == "served"
|
|
|
|
|
|
def test_http_handler_covers_not_found_logging_and_connection_reset():
|
|
mod = load_module()
|
|
|
|
class Collector:
|
|
def render(self):
|
|
return b"metrics\n"
|
|
|
|
class Poller:
|
|
def is_healthy(self):
|
|
return True
|
|
|
|
handler_class = mod.make_handler(Collector(), Poller())
|
|
handler = object.__new__(handler_class)
|
|
handler.path = "/missing"
|
|
statuses = []
|
|
handler.send_error = statuses.append
|
|
handler.do_GET()
|
|
assert statuses == [404]
|
|
assert handler.log_message("ignored") is None
|
|
|
|
class Reset(BytesIO):
|
|
def write(self, _body):
|
|
raise ConnectionResetError
|
|
|
|
handler.path = "/metrics"
|
|
handler.send_response = lambda _status: None
|
|
handler.send_header = lambda _name, _value: None
|
|
handler.end_headers = lambda: None
|
|
handler.wfile = Reset()
|
|
handler.do_GET()
|