Claude (Opus 4.6)
4f80cfcd58
Co-authored-by: Claude (Opus 4.6) <claude@hermes.local> Co-committed-by: Claude (Opus 4.6) <claude@hermes.local>
179 lines
7.1 KiB
Python
179 lines
7.1 KiB
Python
"""Tests for the cloud API budget tracker (issue #882)."""
|
|
|
|
import time
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
from infrastructure.models.budget import (
|
|
BudgetTracker,
|
|
SpendRecord,
|
|
estimate_cost_usd,
|
|
get_budget_tracker,
|
|
)
|
|
|
|
pytestmark = pytest.mark.unit
|
|
|
|
|
|
# ── estimate_cost_usd ─────────────────────────────────────────────────────────
|
|
|
|
|
|
class TestEstimateCostUsd:
|
|
def test_haiku_cheaper_than_sonnet(self):
|
|
haiku_cost = estimate_cost_usd("claude-haiku-4-5", 1000, 1000)
|
|
sonnet_cost = estimate_cost_usd("claude-sonnet-4-5", 1000, 1000)
|
|
assert haiku_cost < sonnet_cost
|
|
|
|
def test_zero_tokens_is_zero_cost(self):
|
|
assert estimate_cost_usd("gpt-4o", 0, 0) == 0.0
|
|
|
|
def test_unknown_model_uses_default(self):
|
|
cost = estimate_cost_usd("some-unknown-model-xyz", 1000, 1000)
|
|
assert cost > 0 # Uses conservative default, not zero
|
|
|
|
def test_versioned_model_name_matches(self):
|
|
# "claude-haiku-4-5-20251001" should match "haiku"
|
|
cost1 = estimate_cost_usd("claude-haiku-4-5-20251001", 1000, 0)
|
|
cost2 = estimate_cost_usd("claude-haiku-4-5", 1000, 0)
|
|
assert cost1 == cost2
|
|
|
|
def test_gpt4o_mini_cheaper_than_gpt4o(self):
|
|
mini = estimate_cost_usd("gpt-4o-mini", 1000, 1000)
|
|
full = estimate_cost_usd("gpt-4o", 1000, 1000)
|
|
assert mini < full
|
|
|
|
def test_returns_float(self):
|
|
assert isinstance(estimate_cost_usd("haiku", 100, 200), float)
|
|
|
|
|
|
# ── BudgetTracker ─────────────────────────────────────────────────────────────
|
|
|
|
|
|
class TestBudgetTrackerInit:
|
|
def test_creates_with_memory_db(self):
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
assert tracker._db_ok is True
|
|
|
|
def test_in_memory_fallback_empty_on_creation(self):
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
assert tracker._in_memory == []
|
|
|
|
def test_bad_path_uses_memory_fallback(self, tmp_path):
|
|
bad_path = str(tmp_path / "nonexistent" / "x" / "budget.db")
|
|
# Should not raise — just log and continue with memory fallback
|
|
# (actually will create parent dirs, so test with truly bad path)
|
|
tracker = BudgetTracker.__new__(BudgetTracker)
|
|
tracker._db_path = bad_path
|
|
tracker._lock = __import__("threading").Lock()
|
|
tracker._in_memory = []
|
|
tracker._db_ok = False
|
|
# Record to in-memory fallback
|
|
tracker._in_memory.append(
|
|
SpendRecord(time.time(), "test", "model", 100, 100, 0.001, "cloud")
|
|
)
|
|
assert len(tracker._in_memory) == 1
|
|
|
|
|
|
class TestBudgetTrackerRecordSpend:
|
|
def test_record_spend_returns_cost(self):
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
cost = tracker.record_spend("anthropic", "claude-haiku-4-5", 100, 200)
|
|
assert cost > 0
|
|
|
|
def test_record_spend_explicit_cost(self):
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
cost = tracker.record_spend("anthropic", "model", cost_usd=1.23)
|
|
assert cost == pytest.approx(1.23)
|
|
|
|
def test_record_spend_accumulates(self):
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
tracker.record_spend("openai", "gpt-4o", cost_usd=0.01)
|
|
tracker.record_spend("openai", "gpt-4o", cost_usd=0.02)
|
|
assert tracker.get_daily_spend() == pytest.approx(0.03, abs=1e-9)
|
|
|
|
def test_record_spend_with_tier_label(self):
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
cost = tracker.record_spend("anthropic", "haiku", tier="cloud_api")
|
|
assert cost >= 0
|
|
|
|
def test_monthly_spend_includes_daily(self):
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
tracker.record_spend("anthropic", "haiku", cost_usd=5.00)
|
|
assert tracker.get_monthly_spend() >= tracker.get_daily_spend()
|
|
|
|
|
|
class TestBudgetTrackerCloudAllowed:
|
|
def test_allowed_when_no_spend(self):
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
with (
|
|
patch.object(type(tracker._get_budget() if hasattr(tracker, "_get_budget") else tracker), "tier_cloud_daily_budget_usd", 5.0, create=True),
|
|
):
|
|
# Settings-based check — use real settings (5.0 default, 0 spent)
|
|
assert tracker.cloud_allowed() is True
|
|
|
|
def test_blocked_when_daily_limit_exceeded(self):
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
tracker.record_spend("anthropic", "haiku", cost_usd=999.0)
|
|
# With default daily limit of 5.0, 999 should block
|
|
assert tracker.cloud_allowed() is False
|
|
|
|
def test_allowed_when_daily_limit_zero(self):
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
tracker.record_spend("anthropic", "haiku", cost_usd=999.0)
|
|
with (
|
|
patch("infrastructure.models.budget.settings") as mock_settings,
|
|
):
|
|
mock_settings.tier_cloud_daily_budget_usd = 0 # disabled
|
|
mock_settings.tier_cloud_monthly_budget_usd = 0 # disabled
|
|
assert tracker.cloud_allowed() is True
|
|
|
|
def test_blocked_when_monthly_limit_exceeded(self):
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
tracker.record_spend("anthropic", "haiku", cost_usd=999.0)
|
|
with patch("infrastructure.models.budget.settings") as mock_settings:
|
|
mock_settings.tier_cloud_daily_budget_usd = 0 # daily disabled
|
|
mock_settings.tier_cloud_monthly_budget_usd = 10.0
|
|
assert tracker.cloud_allowed() is False
|
|
|
|
|
|
class TestBudgetTrackerSummary:
|
|
def test_summary_keys_present(self):
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
summary = tracker.get_summary()
|
|
assert "daily_usd" in summary
|
|
assert "monthly_usd" in summary
|
|
assert "daily_limit_usd" in summary
|
|
assert "monthly_limit_usd" in summary
|
|
assert "daily_ok" in summary
|
|
assert "monthly_ok" in summary
|
|
|
|
def test_summary_daily_ok_true_on_empty(self):
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
summary = tracker.get_summary()
|
|
assert summary["daily_ok"] is True
|
|
assert summary["monthly_ok"] is True
|
|
|
|
def test_summary_daily_ok_false_when_exceeded(self):
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
tracker.record_spend("openai", "gpt-4o", cost_usd=999.0)
|
|
summary = tracker.get_summary()
|
|
assert summary["daily_ok"] is False
|
|
|
|
|
|
# ── Singleton ─────────────────────────────────────────────────────────────────
|
|
|
|
|
|
class TestGetBudgetTrackerSingleton:
|
|
def test_returns_budget_tracker(self):
|
|
import infrastructure.models.budget as bmod
|
|
bmod._budget_tracker = None
|
|
tracker = get_budget_tracker()
|
|
assert isinstance(tracker, BudgetTracker)
|
|
|
|
def test_returns_same_instance(self):
|
|
import infrastructure.models.budget as bmod
|
|
bmod._budget_tracker = None
|
|
t1 = get_budget_tracker()
|
|
t2 = get_budget_tracker()
|
|
assert t1 is t2
|