599 lines
25 KiB
Python
599 lines
25 KiB
Python
"""Unit tests for models/budget.py — comprehensive coverage for budget management.
|
|
|
|
Tests budget allocation, tracking, limit enforcement, and edge cases including:
|
|
- Zero budget scenarios
|
|
- Over-budget handling
|
|
- Budget reset behavior
|
|
- In-memory fallback when DB is unavailable
|
|
"""
|
|
|
|
import threading
|
|
import time
|
|
from datetime import UTC, date, datetime, timedelta
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
from infrastructure.models.budget import (
|
|
BudgetTracker,
|
|
SpendRecord,
|
|
estimate_cost_usd,
|
|
get_budget_tracker,
|
|
)
|
|
|
|
pytestmark = pytest.mark.unit
|
|
|
|
|
|
# ── Test SpendRecord dataclass ────────────────────────────────────────────────
|
|
|
|
|
|
class TestSpendRecord:
|
|
"""Tests for the SpendRecord dataclass."""
|
|
|
|
def test_spend_record_creation(self):
|
|
"""Test creating a SpendRecord with all fields."""
|
|
ts = time.time()
|
|
record = SpendRecord(
|
|
ts=ts,
|
|
provider="anthropic",
|
|
model="claude-haiku-4-5",
|
|
tokens_in=100,
|
|
tokens_out=200,
|
|
cost_usd=0.001,
|
|
tier="cloud",
|
|
)
|
|
assert record.ts == ts
|
|
assert record.provider == "anthropic"
|
|
assert record.model == "claude-haiku-4-5"
|
|
assert record.tokens_in == 100
|
|
assert record.tokens_out == 200
|
|
assert record.cost_usd == 0.001
|
|
assert record.tier == "cloud"
|
|
|
|
def test_spend_record_with_zero_tokens(self):
|
|
"""Test SpendRecord with zero tokens."""
|
|
ts = time.time()
|
|
record = SpendRecord(ts=ts, provider="openai", model="gpt-4o", tokens_in=0, tokens_out=0, cost_usd=0.0, tier="cloud")
|
|
assert record.tokens_in == 0
|
|
assert record.tokens_out == 0
|
|
|
|
|
|
# ── Test estimate_cost_usd function ───────────────────────────────────────────
|
|
|
|
|
|
class TestEstimateCostUsd:
|
|
"""Tests for the estimate_cost_usd function."""
|
|
|
|
def test_haiku_cheaper_than_sonnet(self):
|
|
"""Haiku should be cheaper than Sonnet for same tokens."""
|
|
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):
|
|
"""Zero tokens should result in zero cost."""
|
|
assert estimate_cost_usd("gpt-4o", 0, 0) == 0.0
|
|
|
|
def test_only_input_tokens(self):
|
|
"""Cost calculation with only input tokens."""
|
|
cost = estimate_cost_usd("gpt-4o", 1000, 0)
|
|
expected = (1000 * 0.0025) / 1000.0 # $0.0025 per 1K input tokens
|
|
assert cost == pytest.approx(expected)
|
|
|
|
def test_only_output_tokens(self):
|
|
"""Cost calculation with only output tokens."""
|
|
cost = estimate_cost_usd("gpt-4o", 0, 1000)
|
|
expected = (1000 * 0.01) / 1000.0 # $0.01 per 1K output tokens
|
|
assert cost == pytest.approx(expected)
|
|
|
|
def test_unknown_model_uses_default(self):
|
|
"""Unknown model should use conservative default cost."""
|
|
cost = estimate_cost_usd("some-unknown-model-xyz", 1000, 1000)
|
|
assert cost > 0 # Uses conservative default, not zero
|
|
# Default is 0.003 input, 0.015 output per 1K
|
|
expected = (1000 * 0.003 + 1000 * 0.015) / 1000.0
|
|
assert cost == pytest.approx(expected)
|
|
|
|
def test_versioned_model_name_matches(self):
|
|
"""Versioned model names should match base model rates."""
|
|
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):
|
|
"""GPT-4o mini should be cheaper than GPT-4o."""
|
|
mini = estimate_cost_usd("gpt-4o-mini", 1000, 1000)
|
|
full = estimate_cost_usd("gpt-4o", 1000, 1000)
|
|
assert mini < full
|
|
|
|
def test_opus_most_expensive_claude(self):
|
|
"""Opus should be the most expensive Claude model."""
|
|
opus = estimate_cost_usd("claude-opus-4-5", 1000, 1000)
|
|
sonnet = estimate_cost_usd("claude-sonnet-4-5", 1000, 1000)
|
|
haiku = estimate_cost_usd("claude-haiku-4-5", 1000, 1000)
|
|
assert opus > sonnet > haiku
|
|
|
|
def test_grok_variants(self):
|
|
"""Test Grok model cost estimation."""
|
|
cost = estimate_cost_usd("grok-3", 1000, 1000)
|
|
assert cost > 0
|
|
cost_fast = estimate_cost_usd("grok-3-fast", 1000, 1000)
|
|
assert cost_fast > 0
|
|
|
|
def test_case_insensitive_matching(self):
|
|
"""Model name matching should be case insensitive."""
|
|
cost_lower = estimate_cost_usd("claude-haiku-4-5", 1000, 0)
|
|
cost_upper = estimate_cost_usd("CLAUDE-HAIKU-4-5", 1000, 0)
|
|
cost_mixed = estimate_cost_usd("Claude-Haiku-4-5", 1000, 0)
|
|
assert cost_lower == cost_upper == cost_mixed
|
|
|
|
def test_returns_float(self):
|
|
"""Function should always return a float."""
|
|
assert isinstance(estimate_cost_usd("haiku", 100, 200), float)
|
|
assert isinstance(estimate_cost_usd("unknown-model", 100, 200), float)
|
|
assert isinstance(estimate_cost_usd("haiku", 0, 0), float)
|
|
|
|
|
|
# ── Test BudgetTracker initialization ─────────────────────────────────────────
|
|
|
|
|
|
class TestBudgetTrackerInit:
|
|
"""Tests for BudgetTracker initialization."""
|
|
|
|
def test_creates_with_memory_db(self):
|
|
"""Tracker should initialize with in-memory database."""
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
assert tracker._db_ok is True
|
|
|
|
def test_in_memory_fallback_empty_on_creation(self):
|
|
"""In-memory fallback should start empty."""
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
assert tracker._in_memory == []
|
|
|
|
def test_custom_db_path(self, tmp_path):
|
|
"""Tracker should use custom database path."""
|
|
db_file = tmp_path / "custom_budget.db"
|
|
tracker = BudgetTracker(db_path=str(db_file))
|
|
assert tracker._db_ok is True
|
|
assert tracker._db_path == str(db_file)
|
|
assert db_file.exists()
|
|
|
|
def test_db_path_directory_creation(self, tmp_path):
|
|
"""Tracker should create parent directories if needed."""
|
|
db_file = tmp_path / "nested" / "dirs" / "budget.db"
|
|
tracker = BudgetTracker(db_path=str(db_file))
|
|
assert tracker._db_ok is True
|
|
assert db_file.parent.exists()
|
|
|
|
def test_invalid_db_path_fallback(self):
|
|
"""Tracker should fallback to in-memory on invalid path."""
|
|
# Use a path that cannot be created (e.g., permission denied simulation)
|
|
tracker = BudgetTracker.__new__(BudgetTracker)
|
|
tracker._db_path = "/nonexistent/invalid/path/budget.db"
|
|
tracker._lock = threading.Lock()
|
|
tracker._in_memory = []
|
|
tracker._db_ok = False
|
|
# Should still work with in-memory fallback
|
|
cost = tracker.record_spend("test", "model", cost_usd=0.01)
|
|
assert cost == 0.01
|
|
|
|
|
|
# ── Test BudgetTracker record_spend ───────────────────────────────────────────
|
|
|
|
|
|
class TestBudgetTrackerRecordSpend:
|
|
"""Tests for recording spend events."""
|
|
|
|
def test_record_spend_returns_cost(self):
|
|
"""record_spend should return the calculated cost."""
|
|
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):
|
|
"""record_spend should use explicit cost when provided."""
|
|
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):
|
|
"""Multiple spend records should accumulate correctly."""
|
|
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):
|
|
"""record_spend should accept custom tier labels."""
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
cost = tracker.record_spend("anthropic", "haiku", tier="cloud_api")
|
|
assert cost >= 0
|
|
|
|
def test_record_spend_with_provider(self):
|
|
"""record_spend should track provider correctly."""
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
tracker.record_spend("openai", "gpt-4o", cost_usd=0.01)
|
|
tracker.record_spend("anthropic", "claude-haiku", cost_usd=0.02)
|
|
assert tracker.get_daily_spend() == pytest.approx(0.03, abs=1e-9)
|
|
|
|
def test_record_zero_cost(self):
|
|
"""Recording zero cost should work correctly."""
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
cost = tracker.record_spend("test", "model", cost_usd=0.0)
|
|
assert cost == 0.0
|
|
assert tracker.get_daily_spend() == 0.0
|
|
|
|
def test_record_negative_cost(self):
|
|
"""Recording negative cost (refund) should work."""
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
cost = tracker.record_spend("test", "model", cost_usd=-0.50)
|
|
assert cost == -0.50
|
|
assert tracker.get_daily_spend() == -0.50
|
|
|
|
|
|
# ── Test BudgetTracker daily/monthly spend queries ────────────────────────────
|
|
|
|
|
|
class TestBudgetTrackerSpendQueries:
|
|
"""Tests for daily and monthly spend queries."""
|
|
|
|
def test_monthly_spend_includes_daily(self):
|
|
"""Monthly spend should be >= daily spend."""
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
tracker.record_spend("anthropic", "haiku", cost_usd=5.00)
|
|
assert tracker.get_monthly_spend() >= tracker.get_daily_spend()
|
|
|
|
def test_get_daily_spend_empty(self):
|
|
"""Daily spend should be zero when no records."""
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
assert tracker.get_daily_spend() == 0.0
|
|
|
|
def test_get_monthly_spend_empty(self):
|
|
"""Monthly spend should be zero when no records."""
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
assert tracker.get_monthly_spend() == 0.0
|
|
|
|
def test_daily_spend_isolation(self):
|
|
"""Daily spend should only include today's records, not old ones."""
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
# Force use of in-memory fallback
|
|
tracker._db_ok = False
|
|
|
|
# Add record for today
|
|
today_ts = datetime.combine(date.today(), datetime.min.time(), tzinfo=UTC).timestamp()
|
|
tracker._in_memory.append(
|
|
SpendRecord(today_ts + 3600, "test", "model", 0, 0, 1.0, "cloud")
|
|
)
|
|
|
|
# Add old record (2 days ago)
|
|
old_ts = (datetime.now(UTC) - timedelta(days=2)).timestamp()
|
|
tracker._in_memory.append(
|
|
SpendRecord(old_ts, "test", "old_model", 0, 0, 2.0, "cloud")
|
|
)
|
|
|
|
# Daily should only include today's 1.0
|
|
assert tracker.get_daily_spend() == pytest.approx(1.0, abs=1e-9)
|
|
# Monthly should include both (both are in current month)
|
|
assert tracker.get_monthly_spend() == pytest.approx(3.0, abs=1e-9)
|
|
|
|
|
|
# ── Test BudgetTracker cloud_allowed ──────────────────────────────────────────
|
|
|
|
|
|
class TestBudgetTrackerCloudAllowed:
|
|
"""Tests for cloud budget limit enforcement."""
|
|
|
|
def test_allowed_when_no_spend(self):
|
|
"""Cloud should be allowed when no spend recorded."""
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
assert tracker.cloud_allowed() is True
|
|
|
|
def test_blocked_when_daily_limit_exceeded(self):
|
|
"""Cloud should be blocked when daily limit exceeded."""
|
|
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):
|
|
"""Cloud should be allowed when daily limit is 0 (disabled)."""
|
|
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):
|
|
"""Cloud should be blocked when monthly limit exceeded."""
|
|
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
|
|
|
|
def test_allowed_at_exact_daily_limit(self):
|
|
"""Cloud should be allowed when exactly at daily limit."""
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
with patch("infrastructure.models.budget.settings") as mock_settings:
|
|
mock_settings.tier_cloud_daily_budget_usd = 5.0
|
|
mock_settings.tier_cloud_monthly_budget_usd = 0
|
|
# Record exactly at limit
|
|
tracker.record_spend("test", "model", cost_usd=5.0)
|
|
# At exactly the limit, it should return False (blocked)
|
|
# because spend >= limit
|
|
assert tracker.cloud_allowed() is False
|
|
|
|
def test_allowed_below_daily_limit(self):
|
|
"""Cloud should be allowed when below daily limit."""
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
with patch("infrastructure.models.budget.settings") as mock_settings:
|
|
mock_settings.tier_cloud_daily_budget_usd = 5.0
|
|
mock_settings.tier_cloud_monthly_budget_usd = 0
|
|
tracker.record_spend("test", "model", cost_usd=4.99)
|
|
assert tracker.cloud_allowed() is True
|
|
|
|
def test_zero_budget_blocks_all(self):
|
|
"""Zero budget should block all cloud usage."""
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
with patch("infrastructure.models.budget.settings") as mock_settings:
|
|
mock_settings.tier_cloud_daily_budget_usd = 0.01 # Very small budget
|
|
mock_settings.tier_cloud_monthly_budget_usd = 0
|
|
tracker.record_spend("test", "model", cost_usd=0.02)
|
|
# Over the tiny budget, should be blocked
|
|
assert tracker.cloud_allowed() is False
|
|
|
|
def test_both_limits_checked(self):
|
|
"""Both daily and monthly limits should be checked."""
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
with patch("infrastructure.models.budget.settings") as mock_settings:
|
|
mock_settings.tier_cloud_daily_budget_usd = 100.0
|
|
mock_settings.tier_cloud_monthly_budget_usd = 10.0
|
|
tracker.record_spend("test", "model", cost_usd=15.0)
|
|
# Under daily but over monthly
|
|
assert tracker.cloud_allowed() is False
|
|
|
|
|
|
# ── Test BudgetTracker summary ────────────────────────────────────────────────
|
|
|
|
|
|
class TestBudgetTrackerSummary:
|
|
"""Tests for budget summary functionality."""
|
|
|
|
def test_summary_keys_present(self):
|
|
"""Summary should contain all expected keys."""
|
|
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):
|
|
"""daily_ok and monthly_ok should be True when empty."""
|
|
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):
|
|
"""daily_ok should be False when daily limit exceeded."""
|
|
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
|
|
|
|
def test_summary_monthly_ok_false_when_exceeded(self):
|
|
"""monthly_ok should be False when monthly limit exceeded."""
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
with patch("infrastructure.models.budget.settings") as mock_settings:
|
|
mock_settings.tier_cloud_daily_budget_usd = 0
|
|
mock_settings.tier_cloud_monthly_budget_usd = 10.0
|
|
tracker.record_spend("openai", "gpt-4o", cost_usd=15.0)
|
|
summary = tracker.get_summary()
|
|
assert summary["monthly_ok"] is False
|
|
|
|
def test_summary_values_rounded(self):
|
|
"""Summary values should be rounded appropriately."""
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
tracker.record_spend("test", "model", cost_usd=1.123456789)
|
|
summary = tracker.get_summary()
|
|
# daily_usd should be rounded to 6 decimal places
|
|
assert summary["daily_usd"] == 1.123457
|
|
|
|
def test_summary_with_disabled_limits(self):
|
|
"""Summary should handle disabled limits (0)."""
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
with patch("infrastructure.models.budget.settings") as mock_settings:
|
|
mock_settings.tier_cloud_daily_budget_usd = 0
|
|
mock_settings.tier_cloud_monthly_budget_usd = 0
|
|
tracker.record_spend("test", "model", cost_usd=100.0)
|
|
summary = tracker.get_summary()
|
|
assert summary["daily_limit_usd"] == 0
|
|
assert summary["monthly_limit_usd"] == 0
|
|
assert summary["daily_ok"] is True
|
|
assert summary["monthly_ok"] is True
|
|
|
|
|
|
# ── Test BudgetTracker in-memory fallback ─────────────────────────────────────
|
|
|
|
|
|
class TestBudgetTrackerInMemoryFallback:
|
|
"""Tests for in-memory fallback when DB is unavailable."""
|
|
|
|
def test_in_memory_records_persisted(self):
|
|
"""Records should be stored in memory when DB is unavailable."""
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
# Force DB to appear unavailable
|
|
tracker._db_ok = False
|
|
tracker.record_spend("test", "model", cost_usd=0.01)
|
|
assert len(tracker._in_memory) == 1
|
|
assert tracker._in_memory[0].cost_usd == 0.01
|
|
|
|
def test_in_memory_query_spend(self):
|
|
"""Query spend should work with in-memory fallback."""
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
tracker._db_ok = False
|
|
tracker.record_spend("test", "model", cost_usd=0.01)
|
|
# Query should work from in-memory
|
|
since_ts = (datetime.now(UTC) - timedelta(hours=1)).timestamp()
|
|
result = tracker._query_spend(since_ts)
|
|
assert result == 0.01
|
|
|
|
def test_in_memory_older_records_not_counted(self):
|
|
"""In-memory records older than since_ts should not be counted."""
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
tracker._db_ok = False
|
|
old_ts = (datetime.now(UTC) - timedelta(days=2)).timestamp()
|
|
tracker._in_memory.append(
|
|
SpendRecord(old_ts, "test", "model", 0, 0, 1.0, "cloud")
|
|
)
|
|
# Query for records in last day
|
|
since_ts = (datetime.now(UTC) - timedelta(days=1)).timestamp()
|
|
result = tracker._query_spend(since_ts)
|
|
assert result == 0.0
|
|
|
|
|
|
# ── Test BudgetTracker thread safety ──────────────────────────────────────────
|
|
|
|
|
|
class TestBudgetTrackerThreadSafety:
|
|
"""Tests for thread-safe operations."""
|
|
|
|
def test_concurrent_record_spend(self):
|
|
"""Multiple threads should safely record spend concurrently."""
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
results = []
|
|
errors = []
|
|
|
|
def record_spends():
|
|
try:
|
|
for _ in range(10):
|
|
cost = tracker.record_spend("test", "model", cost_usd=0.01)
|
|
results.append(cost)
|
|
except Exception as e:
|
|
errors.append(e)
|
|
|
|
threads = [threading.Thread(target=record_spends) for _ in range(5)]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
assert len(errors) == 0
|
|
assert len(results) == 50
|
|
assert tracker.get_daily_spend() == pytest.approx(0.50, abs=1e-9)
|
|
|
|
|
|
# ── Test BudgetTracker edge cases ─────────────────────────────────────────────
|
|
|
|
|
|
class TestBudgetTrackerEdgeCases:
|
|
"""Tests for edge cases and boundary conditions."""
|
|
|
|
def test_very_small_cost(self):
|
|
"""Tracker should handle very small costs."""
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
tracker.record_spend("test", "model", cost_usd=0.000001)
|
|
assert tracker.get_daily_spend() == pytest.approx(0.000001, abs=1e-9)
|
|
|
|
def test_very_large_cost(self):
|
|
"""Tracker should handle very large costs."""
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
tracker.record_spend("test", "model", cost_usd=1_000_000.0)
|
|
assert tracker.get_daily_spend() == pytest.approx(1_000_000.0, abs=1e-9)
|
|
|
|
def test_many_records(self):
|
|
"""Tracker should handle many records efficiently."""
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
for i in range(100):
|
|
tracker.record_spend(f"provider_{i}", f"model_{i}", cost_usd=0.01)
|
|
assert tracker.get_daily_spend() == pytest.approx(1.0, abs=1e-9)
|
|
|
|
def test_empty_provider_name(self):
|
|
"""Tracker should handle empty provider name."""
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
cost = tracker.record_spend("", "model", cost_usd=0.01)
|
|
assert cost == 0.01
|
|
|
|
def test_empty_model_name(self):
|
|
"""Tracker should handle empty model name."""
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
cost = tracker.record_spend("provider", "", cost_usd=0.01)
|
|
assert cost == 0.01
|
|
|
|
|
|
# ── Test get_budget_tracker singleton ─────────────────────────────────────────
|
|
|
|
|
|
class TestGetBudgetTrackerSingleton:
|
|
"""Tests for the module-level BudgetTracker singleton."""
|
|
|
|
def test_returns_budget_tracker(self):
|
|
"""Singleton should return a BudgetTracker instance."""
|
|
import infrastructure.models.budget as bmod
|
|
|
|
bmod._budget_tracker = None
|
|
tracker = get_budget_tracker()
|
|
assert isinstance(tracker, BudgetTracker)
|
|
|
|
def test_returns_same_instance(self):
|
|
"""Singleton should return the same instance."""
|
|
import infrastructure.models.budget as bmod
|
|
|
|
bmod._budget_tracker = None
|
|
t1 = get_budget_tracker()
|
|
t2 = get_budget_tracker()
|
|
assert t1 is t2
|
|
|
|
def test_singleton_persists_state(self):
|
|
"""Singleton should persist state across calls."""
|
|
import infrastructure.models.budget as bmod
|
|
|
|
bmod._budget_tracker = None
|
|
tracker1 = get_budget_tracker()
|
|
# Record spend
|
|
tracker1.record_spend("test", "model", cost_usd=1.0)
|
|
# Get singleton again
|
|
tracker2 = get_budget_tracker()
|
|
assert tracker1 is tracker2
|
|
|
|
|
|
# ── Test BudgetTracker with mocked settings ───────────────────────────────────
|
|
|
|
|
|
class TestBudgetTrackerWithMockedSettings:
|
|
"""Tests using mocked settings for different scenarios."""
|
|
|
|
def test_high_daily_limit(self):
|
|
"""Test with high daily limit."""
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
with patch("infrastructure.models.budget.settings") as mock_settings:
|
|
mock_settings.tier_cloud_daily_budget_usd = 1000.0
|
|
mock_settings.tier_cloud_monthly_budget_usd = 10000.0
|
|
tracker.record_spend("test", "model", cost_usd=500.0)
|
|
assert tracker.cloud_allowed() is True
|
|
|
|
def test_low_daily_limit(self):
|
|
"""Test with low daily limit."""
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
with patch("infrastructure.models.budget.settings") as mock_settings:
|
|
mock_settings.tier_cloud_daily_budget_usd = 1.0
|
|
mock_settings.tier_cloud_monthly_budget_usd = 100.0
|
|
tracker.record_spend("test", "model", cost_usd=2.0)
|
|
assert tracker.cloud_allowed() is False
|
|
|
|
def test_only_monthly_limit_enabled(self):
|
|
"""Test with only monthly limit enabled."""
|
|
tracker = BudgetTracker(db_path=":memory:")
|
|
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 = 50.0
|
|
tracker.record_spend("test", "model", cost_usd=30.0)
|
|
assert tracker.cloud_allowed() is True
|
|
tracker.record_spend("test", "model", cost_usd=25.0)
|
|
assert tracker.cloud_allowed() is False
|