Timmy-time-dashboard/tests/infrastructure/test_budget_tracker.py

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