276 lines
8.2 KiB
Python
276 lines
8.2 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_provider_queries_validate_credentials_and_response_shape(
|
|
tmp_path, monkeypatch
|
|
):
|
|
mod = load_module()
|
|
credentials = tmp_path / "credentials.json"
|
|
credentials.write_text("{}", encoding="utf-8")
|
|
monkeypatch.setattr(mod, "CLAUDE_CREDENTIALS", credentials)
|
|
|
|
with pytest.raises(RuntimeError, match="credentials are incomplete"):
|
|
mod.query_claude()
|
|
|
|
credentials.write_text(
|
|
json.dumps({"claudeAiOauth": {"accessToken": "runtime-only-token"}}),
|
|
encoding="utf-8",
|
|
)
|
|
bodies = iter([{"five_hour": {}}, []])
|
|
requests = []
|
|
|
|
class Response:
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *_args):
|
|
return None
|
|
|
|
def read(self, maximum):
|
|
assert maximum == 1 << 20
|
|
return json.dumps(next(bodies)).encode()
|
|
|
|
def open_request(request, timeout):
|
|
requests.append(request)
|
|
assert timeout == 15
|
|
return Response()
|
|
|
|
monkeypatch.setattr(mod, "urlopen", open_request)
|
|
assert mod.query_claude() == {"five_hour": {}}
|
|
assert requests[0].get_header("Authorization") == "Bearer runtime-only-token"
|
|
with pytest.raises(RuntimeError, match="response is not an object"):
|
|
mod.query_claude()
|
|
|
|
assert mod.query_codex is mod.codex_query.query_codex
|
|
|
|
|
|
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"].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()
|