"""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