📦 EqualifyEverything / equalify-reflow

📄 test_llm_cost.py · 189 lines
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
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189"""Tests for centralized LLM cost configuration module."""

import pytest
from src.shared.llm_cost import (
    DEFAULT_PRICING,
    LLMPricing,
    calculate_estimated_cost,
    format_cost_dollars,
)


@pytest.mark.unit
class TestLLMPricing:
    """Test LLMPricing dataclass."""

    def test_llm_pricing_creation(self):
        """Test creating LLMPricing with all fields."""
        pricing = LLMPricing(
            input_cost_per_token_cents=0.0001,
            output_cost_per_token_cents=0.0005,
            model_name="Test Model",
        )

        assert pricing.input_cost_per_token_cents == 0.0001
        assert pricing.output_cost_per_token_cents == 0.0005
        assert pricing.model_name == "Test Model"

    def test_llm_pricing_with_default_model_name(self):
        """Test LLMPricing defaults to 'Unknown Model' if not provided."""
        pricing = LLMPricing(
            input_cost_per_token_cents=0.0001,
            output_cost_per_token_cents=0.0005,
        )

        assert pricing.model_name == "Unknown Model"


@pytest.mark.unit
class TestDefaultPricing:
    """Test DEFAULT_PRICING constant."""

    def test_default_pricing_values(self):
        """Test DEFAULT_PRICING has correct Claude Haiku 4.5 pricing."""
        # Claude Haiku 4.5 pricing: $1.00/1M input, $5.00/1M output
        assert DEFAULT_PRICING.input_cost_per_token_cents == 0.0001
        assert DEFAULT_PRICING.output_cost_per_token_cents == 0.0005
        assert "Claude Haiku 4.5" in DEFAULT_PRICING.model_name
        assert "Bedrock" in DEFAULT_PRICING.model_name


@pytest.mark.unit
class TestCalculateEstimatedCost:
    """Test calculate_estimated_cost function."""

    def test_calculate_cost_basic(self):
        """Test basic cost calculation with default pricing."""
        cost = calculate_estimated_cost(input_tokens=1000, output_tokens=100)

        # Expected: 1000 * 0.0001 + 100 * 0.0005 = 0.1 + 0.05 = 0.15 cents
        assert cost == pytest.approx(0.15)

    def test_calculate_cost_zero_tokens(self):
        """Test cost calculation with zero tokens."""
        cost = calculate_estimated_cost(input_tokens=0, output_tokens=0)
        assert cost == 0.0

    def test_calculate_cost_input_only(self):
        """Test cost calculation with only input tokens."""
        cost = calculate_estimated_cost(input_tokens=1000, output_tokens=0)

        # Expected: 1000 * 0.0001 = 0.1 cents
        assert cost == pytest.approx(0.1)

    def test_calculate_cost_output_only(self):
        """Test cost calculation with only output tokens."""
        cost = calculate_estimated_cost(input_tokens=0, output_tokens=100)

        # Expected: 100 * 0.0005 = 0.05 cents
        assert cost == pytest.approx(0.05)

    def test_calculate_cost_large_tokens(self):
        """Test cost calculation with 1 million tokens each."""
        cost = calculate_estimated_cost(input_tokens=1_000_000, output_tokens=1_000_000)

        # Expected: 1M * 0.0001 + 1M * 0.0005 = 100 + 500 = 600 cents = $6.00
        assert cost == pytest.approx(600.0)

    def test_calculate_cost_custom_pricing(self):
        """Test cost calculation with custom pricing."""
        custom_pricing = LLMPricing(
            input_cost_per_token_cents=0.0002,
            output_cost_per_token_cents=0.001,
            model_name="Custom Model",
        )

        cost = calculate_estimated_cost(
            input_tokens=1000, output_tokens=100, pricing=custom_pricing
        )

        # Expected: 1000 * 0.0002 + 100 * 0.001 = 0.2 + 0.1 = 0.3 cents
        assert cost == pytest.approx(0.3)

    def test_calculate_cost_matches_example_in_docstring(self):
        """Test that cost calculation matches the docstring example."""
        # From docstring: 1000 input, 100 output should be $0.001500
        cost = calculate_estimated_cost(1000, 100)
        dollars = cost / 100

        assert dollars == pytest.approx(0.0015)


@pytest.mark.unit
class TestFormatCostDollars:
    """Test format_cost_dollars function."""

    def test_format_very_small_cost(self):
        """Test formatting very small cost (< $0.01) with high precision."""
        # 0.15 cents = $0.0015
        formatted = format_cost_dollars(0.15)
        assert formatted == "$0.001500"

    def test_format_small_cost(self):
        """Test formatting small cost (< $0.01) with high precision."""
        # 0.5 cents = $0.005
        formatted = format_cost_dollars(0.5)
        assert formatted == "$0.005000"

    def test_format_moderate_cost(self):
        """Test formatting moderate cost (>= $0.01) with 2 decimal places."""
        # 150 cents = $1.50
        formatted = format_cost_dollars(150.0)
        assert formatted == "$1.50"

    def test_format_large_cost(self):
        """Test formatting large cost with 2 decimal places."""
        # 12345 cents = $123.45
        formatted = format_cost_dollars(12345.0)
        assert formatted == "$123.45"

    def test_format_zero_cost(self):
        """Test formatting zero cost."""
        formatted = format_cost_dollars(0.0)
        assert formatted == "$0.000000"

    def test_format_one_cent_boundary(self):
        """Test formatting at the $0.01 boundary."""
        # Just below $0.01 (0.99 cents)
        formatted_below = format_cost_dollars(0.99)
        assert formatted_below == "$0.009900"

        # Exactly $0.01 (1.0 cents)
        formatted_exact = format_cost_dollars(1.0)
        assert formatted_exact == "$0.01"

        # Just above $0.01 (1.01 cents)
        formatted_above = format_cost_dollars(1.01)
        assert formatted_above == "$0.01"


@pytest.mark.unit
class TestIntegration:
    """Integration tests combining multiple functions."""

    def test_calculate_and_format_workflow(self):
        """Test typical workflow: calculate cost then format it."""
        # Calculate cost for 10,000 input tokens and 1,000 output tokens
        cost_cents = calculate_estimated_cost(10_000, 1_000)

        # Expected: 10000 * 0.0001 + 1000 * 0.0005 = 1.0 + 0.5 = 1.5 cents
        assert cost_cents == pytest.approx(1.5)

        # Format the cost (1.5 cents = $0.015, which is >= $0.01, so formatted with 2 decimals)
        formatted = format_cost_dollars(cost_cents)
        assert formatted == "$0.01"

    def test_cost_calculation_consistency_with_base_agent(self):
        """Test that cost calculation is consistent with what agents expect."""
        # This mirrors the test in test_base_agent.py
        input_tokens = 1000
        output_tokens = 100

        cost = calculate_estimated_cost(input_tokens, output_tokens)

        expected_input_cost = input_tokens * DEFAULT_PRICING.input_cost_per_token_cents
        expected_output_cost = output_tokens * DEFAULT_PRICING.output_cost_per_token_cents
        expected_total = expected_input_cost + expected_output_cost

        assert cost == pytest.approx(expected_total)