PlanOpticon

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

Keyboard Shortcuts

Open search /
Next entry (timeline) j
Previous entry (timeline) k
Open focused entry Enter
Show this help ?
Toggle theme Top nav button