Spaces:
Runtime error
Runtime error
File size: 4,000 Bytes
790b5af | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | """Unit tests for cost tracking utility."""
import pytest
from src.utils.cost_tracker import CostTracker, TokenUsage, BudgetExceededError
def test_token_usage():
"""Test TokenUsage dataclass."""
usage = TokenUsage(input_tokens=1000, output_tokens=500, model="test-model")
assert usage.input_tokens == 1000
assert usage.output_tokens == 500
assert usage.total_tokens == 1500
assert usage.model == "test-model"
def test_calculate_cost_claude():
"""Test cost calculation for Claude model."""
tracker = CostTracker()
# Claude Sonnet 4.5: $3/1M input, $15/1M output
cost = tracker.calculate_cost(
"anthropic/claude-sonnet-4.5", input_tokens=10_000, output_tokens=5_000
)
expected = (10_000 * 3.00 / 1_000_000) + (5_000 * 15.00 / 1_000_000)
assert abs(cost - expected) < 0.0001
def test_calculate_cost_gpt5_mini():
"""Test cost calculation for GPT-5-mini."""
tracker = CostTracker()
# GPT-5-mini: $0.25/1M input, $2.00/1M output
cost = tracker.calculate_cost(
"openai/gpt-5-mini", input_tokens=100_000, output_tokens=50_000
)
expected = (100_000 * 0.25 / 1_000_000) + (50_000 * 2.00 / 1_000_000)
assert abs(cost - expected) < 0.0001
def test_calculate_cost_ollama():
"""Test cost calculation for Ollama (should be free)."""
tracker = CostTracker()
cost = tracker.calculate_cost("ollama", input_tokens=100_000, output_tokens=50_000)
assert cost == 0.0
def test_track_usage():
"""Test usage tracking updates total cost."""
tracker = CostTracker()
assert tracker.total_cost == 0.0
assert len(tracker.usage_history) == 0
# Track first usage
cost1 = tracker.track_usage("openai/gpt-5-mini", 10_000, 5_000)
assert tracker.total_cost == cost1
assert len(tracker.usage_history) == 1
# Track second usage
cost2 = tracker.track_usage("anthropic/claude-sonnet-4.5", 20_000, 10_000)
assert abs(tracker.total_cost - (cost1 + cost2)) < 0.0001
assert len(tracker.usage_history) == 2
def test_budget_check_within_limit():
"""Test budget check passes when within limit."""
tracker = CostTracker()
tracker.track_usage("openai/gpt-5-mini", 10_000, 5_000)
# Should not raise
tracker.check_budget(max_budget=1.0)
def test_budget_check_exceeds_limit():
"""Test budget check raises exception when exceeded."""
tracker = CostTracker()
# Track expensive usage
tracker.track_usage("anthropic/claude-3.5-sonnet", 100_000, 200_000)
# Should raise BudgetExceededError
with pytest.raises(BudgetExceededError, match="exceeds budget"):
tracker.check_budget(max_budget=0.01)
def test_get_summary():
"""Test cost summary generation."""
tracker = CostTracker()
tracker.track_usage("openai/gpt-5-mini", 10_000, 5_000)
tracker.track_usage("anthropic/claude-sonnet-4.5", 20_000, 10_000)
tracker.track_usage("openai/gpt-5-mini", 5_000, 2_500) # Same model again
summary = tracker.get_summary()
assert summary["calls"] == 3
assert summary["total_input_tokens"] == 35_000
assert summary["total_output_tokens"] == 17_500
assert summary["total_tokens"] == 52_500
assert "openai/gpt-5-mini" in summary["by_model"]
assert "anthropic/claude-sonnet-4.5" in summary["by_model"]
# Check model-specific costs
gpt_mini = summary["by_model"]["openai/gpt-5-mini"]
assert gpt_mini["input_tokens"] == 15_000 # 10k + 5k
assert gpt_mini["output_tokens"] == 7_500 # 5k + 2.5k
claude = summary["by_model"]["anthropic/claude-sonnet-4.5"]
assert claude["input_tokens"] == 20_000
assert claude["output_tokens"] == 10_000
def test_unknown_model_fallback():
"""Test unknown model falls back to GPT-5-mini pricing."""
tracker = CostTracker()
cost = tracker.calculate_cost("unknown/model", 10_000, 5_000)
expected_cost = tracker.calculate_cost("openai/gpt-5-mini", 10_000, 5_000)
assert abs(cost - expected_cost) < 0.0001
|