|
1
|
"""Tests for the UsageTracker class.""" |
|
2
|
|
|
3
|
import time |
|
4
|
|
|
5
|
from video_processor.utils.usage_tracker import ModelUsage, StepTiming, UsageTracker, _fmt_duration |
|
6
|
|
|
7
|
|
|
8
|
class TestModelUsage: |
|
9
|
def test_total_tokens(self): |
|
10
|
mu = ModelUsage(provider="openai", model="gpt-4o", input_tokens=100, output_tokens=50) |
|
11
|
assert mu.total_tokens == 150 |
|
12
|
|
|
13
|
def test_estimated_cost_known_model(self): |
|
14
|
mu = ModelUsage( |
|
15
|
provider="openai", |
|
16
|
model="gpt-4o", |
|
17
|
input_tokens=1_000_000, |
|
18
|
output_tokens=500_000, |
|
19
|
) |
|
20
|
# gpt-4o: input $2.50/M, output $10.00/M |
|
21
|
expected = 1_000_000 * 2.50 / 1_000_000 + 500_000 * 10.00 / 1_000_000 |
|
22
|
assert abs(mu.estimated_cost - expected) < 0.001 |
|
23
|
|
|
24
|
def test_estimated_cost_unknown_model(self): |
|
25
|
mu = ModelUsage( |
|
26
|
provider="local", |
|
27
|
model="my-custom-model", |
|
28
|
input_tokens=1000, |
|
29
|
output_tokens=500, |
|
30
|
) |
|
31
|
assert mu.estimated_cost == 0.0 |
|
32
|
|
|
33
|
def test_estimated_cost_whisper(self): |
|
34
|
mu = ModelUsage( |
|
35
|
provider="openai", |
|
36
|
model="whisper-1", |
|
37
|
audio_minutes=10.0, |
|
38
|
) |
|
39
|
# whisper-1: $0.006/min |
|
40
|
assert abs(mu.estimated_cost - 0.06) < 0.001 |
|
41
|
|
|
42
|
def test_estimated_cost_partial_match(self): |
|
43
|
mu = ModelUsage( |
|
44
|
provider="openai", |
|
45
|
model="gpt-4o-2024-08-06", |
|
46
|
input_tokens=1_000_000, |
|
47
|
output_tokens=0, |
|
48
|
) |
|
49
|
# Should partial-match to gpt-4o |
|
50
|
assert mu.estimated_cost > 0 |
|
51
|
|
|
52
|
def test_calls_default_zero(self): |
|
53
|
mu = ModelUsage() |
|
54
|
assert mu.calls == 0 |
|
55
|
assert mu.total_tokens == 0 |
|
56
|
assert mu.estimated_cost == 0.0 |
|
57
|
|
|
58
|
|
|
59
|
class TestStepTiming: |
|
60
|
def test_duration_with_times(self): |
|
61
|
st = StepTiming(name="test", start_time=100.0, end_time=105.5) |
|
62
|
assert abs(st.duration - 5.5) < 0.001 |
|
63
|
|
|
64
|
def test_duration_no_end_time(self): |
|
65
|
st = StepTiming(name="test", start_time=100.0) |
|
66
|
assert st.duration == 0.0 |
|
67
|
|
|
68
|
def test_duration_no_start_time(self): |
|
69
|
st = StepTiming(name="test") |
|
70
|
assert st.duration == 0.0 |
|
71
|
|
|
72
|
|
|
73
|
class TestUsageTracker: |
|
74
|
def test_record_single_call(self): |
|
75
|
tracker = UsageTracker() |
|
76
|
tracker.record("openai", "gpt-4o", input_tokens=500, output_tokens=200) |
|
77
|
assert tracker.total_api_calls == 1 |
|
78
|
assert tracker.total_input_tokens == 500 |
|
79
|
assert tracker.total_output_tokens == 200 |
|
80
|
assert tracker.total_tokens == 700 |
|
81
|
|
|
82
|
def test_record_multiple_calls_same_model(self): |
|
83
|
tracker = UsageTracker() |
|
84
|
tracker.record("openai", "gpt-4o", input_tokens=100, output_tokens=50) |
|
85
|
tracker.record("openai", "gpt-4o", input_tokens=200, output_tokens=100) |
|
86
|
assert tracker.total_api_calls == 2 |
|
87
|
assert tracker.total_input_tokens == 300 |
|
88
|
assert tracker.total_output_tokens == 150 |
|
89
|
|
|
90
|
def test_record_multiple_models(self): |
|
91
|
tracker = UsageTracker() |
|
92
|
tracker.record("openai", "gpt-4o", input_tokens=100, output_tokens=50) |
|
93
|
tracker.record( |
|
94
|
"anthropic", "claude-sonnet-4-5-20250929", input_tokens=200, output_tokens=100 |
|
95
|
) |
|
96
|
assert tracker.total_api_calls == 2 |
|
97
|
assert tracker.total_input_tokens == 300 |
|
98
|
assert len(tracker._models) == 2 |
|
99
|
|
|
100
|
def test_total_cost(self): |
|
101
|
tracker = UsageTracker() |
|
102
|
tracker.record("openai", "gpt-4o", input_tokens=1_000_000, output_tokens=500_000) |
|
103
|
cost = tracker.total_cost |
|
104
|
assert cost > 0 |
|
105
|
|
|
106
|
def test_start_and_end_step(self): |
|
107
|
tracker = UsageTracker() |
|
108
|
tracker.start_step("Frame extraction") |
|
109
|
time.sleep(0.01) |
|
110
|
tracker.end_step() |
|
111
|
|
|
112
|
assert len(tracker._steps) == 1 |
|
113
|
assert tracker._steps[0].name == "Frame extraction" |
|
114
|
assert tracker._steps[0].duration > 0 |
|
115
|
|
|
116
|
def test_start_step_auto_closes_previous(self): |
|
117
|
tracker = UsageTracker() |
|
118
|
tracker.start_step("Step 1") |
|
119
|
time.sleep(0.01) |
|
120
|
tracker.start_step("Step 2") |
|
121
|
# Step 1 should have been auto-closed |
|
122
|
assert len(tracker._steps) == 1 |
|
123
|
assert tracker._steps[0].name == "Step 1" |
|
124
|
assert tracker._steps[0].duration > 0 |
|
125
|
# Step 2 is current |
|
126
|
assert tracker._current_step.name == "Step 2" |
|
127
|
|
|
128
|
def test_end_step_when_none(self): |
|
129
|
tracker = UsageTracker() |
|
130
|
tracker.end_step() # Should not raise |
|
131
|
assert len(tracker._steps) == 0 |
|
132
|
|
|
133
|
def test_total_duration(self): |
|
134
|
tracker = UsageTracker() |
|
135
|
time.sleep(0.01) |
|
136
|
assert tracker.total_duration > 0 |
|
137
|
|
|
138
|
def test_format_summary_empty(self): |
|
139
|
tracker = UsageTracker() |
|
140
|
summary = tracker.format_summary() |
|
141
|
assert "PROCESSING SUMMARY" in summary |
|
142
|
assert "Total time" in summary |
|
143
|
|
|
144
|
def test_format_summary_with_usage(self): |
|
145
|
tracker = UsageTracker() |
|
146
|
tracker.record("openai", "gpt-4o", input_tokens=1000, output_tokens=500) |
|
147
|
tracker.start_step("Analysis") |
|
148
|
tracker.end_step() |
|
149
|
|
|
150
|
summary = tracker.format_summary() |
|
151
|
assert "API Calls" in summary |
|
152
|
assert "Tokens" in summary |
|
153
|
assert "gpt-4o" in summary |
|
154
|
assert "Analysis" in summary |
|
155
|
|
|
156
|
def test_format_summary_with_audio(self): |
|
157
|
tracker = UsageTracker() |
|
158
|
tracker.record("openai", "whisper-1", audio_minutes=5.0) |
|
159
|
summary = tracker.format_summary() |
|
160
|
assert "whisper" in summary |
|
161
|
assert "5.0m" in summary |
|
162
|
|
|
163
|
def test_format_summary_cost_display(self): |
|
164
|
tracker = UsageTracker() |
|
165
|
tracker.record("openai", "gpt-4o", input_tokens=1_000_000, output_tokens=500_000) |
|
166
|
summary = tracker.format_summary() |
|
167
|
assert "Estimated total cost: $" in summary |
|
168
|
|
|
169
|
def test_format_summary_step_percentages(self): |
|
170
|
tracker = UsageTracker() |
|
171
|
# Manually create steps with known timings |
|
172
|
tracker._steps = [ |
|
173
|
StepTiming(name="Step A", start_time=0.0, end_time=1.0), |
|
174
|
StepTiming(name="Step B", start_time=1.0, end_time=3.0), |
|
175
|
] |
|
176
|
summary = tracker.format_summary() |
|
177
|
assert "Step A" in summary |
|
178
|
assert "Step B" in summary |
|
179
|
assert "%" in summary |
|
180
|
|
|
181
|
|
|
182
|
class TestFmtDuration: |
|
183
|
def test_seconds(self): |
|
184
|
assert _fmt_duration(5.3) == "5.3s" |
|
185
|
|
|
186
|
def test_minutes(self): |
|
187
|
result = _fmt_duration(90.0) |
|
188
|
assert result == "1m 30s" |
|
189
|
|
|
190
|
def test_hours(self): |
|
191
|
result = _fmt_duration(3661.0) |
|
192
|
assert result == "1h 1m 1s" |
|
193
|
|
|
194
|
def test_zero(self): |
|
195
|
assert _fmt_duration(0.0) == "0.0s" |
|
196
|
|
|
197
|
def test_just_under_minute(self): |
|
198
|
assert _fmt_duration(59.9) == "59.9s" |
|
199
|
|