466683e14d
Add two tools from the March 23 operational briefing: - src/infrastructure/claude_quota.py: SQLite-backed tracker for Claude API usage (tokens, cost, calls) per day/month. Exposes current_mode() which returns BURST / ACTIVE / RESTING based on daily spend thresholds, enabling the orchestrator to route inference requests according to the metabolic protocol (issue #972). - scripts/claude_quota_check.sh: CLI wrapper with --mode (print mode only) and --json (machine-readable) flags for quick quota inspection from the shell or CI scripts. - tests/infrastructure/test_claude_quota.py: 19 unit tests covering cost calculation, mode thresholds, store CRUD, and convenience functions. Refs #1074
140 lines
5.0 KiB
Python
140 lines
5.0 KiB
Python
"""Tests for the Claude quota tracker and metabolic mode advisor.
|
|
|
|
Refs: #1074
|
|
"""
|
|
|
|
import pytest
|
|
|
|
from infrastructure.claude_quota import (
|
|
ACTIVE_THRESHOLD,
|
|
BURST_THRESHOLD,
|
|
ClaudeCall,
|
|
ClaudeQuotaStore,
|
|
MetabolicMode,
|
|
_mode_for_cost,
|
|
current_mode,
|
|
quota_report,
|
|
record_usage,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def store(tmp_path):
|
|
"""Fresh quota store backed by a temp DB."""
|
|
return ClaudeQuotaStore(db_path=tmp_path / "test_quota.db")
|
|
|
|
|
|
# ── Unit: cost calculation ────────────────────────────────────────────────────
|
|
|
|
|
|
class TestClaudeCallCost:
|
|
def test_haiku_cost(self):
|
|
call = ClaudeCall(model="haiku", input_tokens=1_000_000, output_tokens=0)
|
|
assert call.cost_usd == pytest.approx(0.25)
|
|
|
|
def test_sonnet_output_cost(self):
|
|
call = ClaudeCall(model="sonnet", input_tokens=0, output_tokens=1_000_000)
|
|
assert call.cost_usd == pytest.approx(15.00)
|
|
|
|
def test_opus_combined_cost(self):
|
|
call = ClaudeCall(model="opus", input_tokens=100_000, output_tokens=50_000)
|
|
# input: 100k * 15/1M = 1.50, output: 50k * 75/1M = 3.75 → 5.25
|
|
assert call.cost_usd == pytest.approx(5.25)
|
|
|
|
def test_unknown_model_uses_default(self):
|
|
call = ClaudeCall(model="unknown-model-xyz", input_tokens=1_000_000, output_tokens=0)
|
|
assert call.cost_usd == pytest.approx(3.00) # default input cost
|
|
|
|
def test_zero_tokens_zero_cost(self):
|
|
call = ClaudeCall(model="haiku", input_tokens=0, output_tokens=0)
|
|
assert call.cost_usd == 0.0
|
|
|
|
|
|
# ── Unit: metabolic mode thresholds ──────────────────────────────────────────
|
|
|
|
|
|
class TestMetabolicMode:
|
|
def test_under_burst_threshold(self):
|
|
assert _mode_for_cost(0.0) == "BURST"
|
|
assert _mode_for_cost(BURST_THRESHOLD - 0.01) == "BURST"
|
|
|
|
def test_at_burst_threshold_is_active(self):
|
|
assert _mode_for_cost(BURST_THRESHOLD) == "ACTIVE"
|
|
|
|
def test_between_thresholds(self):
|
|
mid = (BURST_THRESHOLD + ACTIVE_THRESHOLD) / 2
|
|
assert _mode_for_cost(mid) == "ACTIVE"
|
|
|
|
def test_at_active_threshold_is_resting(self):
|
|
assert _mode_for_cost(ACTIVE_THRESHOLD) == "RESTING"
|
|
|
|
def test_over_active_threshold(self):
|
|
assert _mode_for_cost(ACTIVE_THRESHOLD + 10) == "RESTING"
|
|
|
|
|
|
# ── Store: record and query ───────────────────────────────────────────────────
|
|
|
|
|
|
class TestClaudeQuotaStore:
|
|
def test_record_call(self, store):
|
|
call = ClaudeCall(model="haiku", input_tokens=1000, output_tokens=500)
|
|
store.record_call(call)
|
|
summary = store.today_summary()
|
|
assert summary.calls == 1
|
|
assert summary.input_tokens == 1000
|
|
assert summary.output_tokens == 500
|
|
assert summary.cost_usd > 0
|
|
|
|
def test_today_summary_empty_db(self, store):
|
|
summary = store.today_summary()
|
|
assert summary.calls == 0
|
|
assert summary.cost_usd == 0.0
|
|
assert summary.mode == "BURST"
|
|
|
|
def test_month_summary_aggregates_multiple_calls(self, store):
|
|
for _ in range(5):
|
|
store.record_call(ClaudeCall(model="haiku", input_tokens=100, output_tokens=50))
|
|
month = store.month_summary()
|
|
assert month.calls == 5
|
|
assert month.input_tokens == 500
|
|
assert month.output_tokens == 250
|
|
|
|
def test_current_mode_burst_when_empty(self, store):
|
|
assert store.current_mode() == "BURST"
|
|
|
|
def test_current_mode_resting_when_expensive(self, store):
|
|
# Record enough usage to push past ACTIVE_THRESHOLD
|
|
# ACTIVE_THRESHOLD = 5.00, opus input = 15/1M
|
|
# Need >5.00: 5.00/15 * 1M ≈ 333_334 input tokens
|
|
store.record_call(
|
|
ClaudeCall(model="opus", input_tokens=400_000, output_tokens=0)
|
|
)
|
|
mode = store.current_mode()
|
|
assert mode == "RESTING"
|
|
|
|
def test_summary_as_dict(self, store):
|
|
summary = store.today_summary()
|
|
d = summary.as_dict()
|
|
assert "period" in d
|
|
assert "calls" in d
|
|
assert "cost_usd" in d
|
|
assert "mode" in d
|
|
|
|
|
|
# ── Convenience functions ─────────────────────────────────────────────────────
|
|
|
|
|
|
class TestConvenienceFunctions:
|
|
def test_record_usage_does_not_raise(self):
|
|
# Uses module-level store; should not raise even if DB path issues
|
|
record_usage(model="haiku", input_tokens=10, output_tokens=5, task_label="test")
|
|
|
|
def test_current_mode_returns_valid_mode(self):
|
|
mode = current_mode()
|
|
assert mode in ("BURST", "ACTIVE", "RESTING")
|
|
|
|
def test_quota_report_returns_string(self):
|
|
report = quota_report()
|
|
assert isinstance(report, str)
|
|
assert "BURST" in report or "ACTIVE" in report or "RESTING" in report
|