@@ -11,73 +11,67 @@ class TestAI:
1111 async def test_llm_basic_call (self , mock_litellm ):
1212 """Test basic LLM call functionality"""
1313 result = await llm ("gpt-4" , "Test prompt" , 100 )
14-
14+
1515 # Verify litellm was called correctly
1616 mock_litellm .assert_called_once_with (
17- "gpt-4" ,
18- [{"role" : "user" , "content" : "Test prompt" }],
19- max_tokens = 100
17+ "gpt-4" , [{"role" : "user" , "content" : "Test prompt" }], max_tokens = 100
2018 )
21-
19+
2220 # Verify response is returned
2321 assert result == mock_litellm .return_value
24-
22+
2523 @pytest .mark .asyncio
2624 async def test_llm_without_max_tokens (self , mock_litellm ):
2725 """Test LLM call without specifying max tokens"""
2826 await llm ("gpt-4" , "Test prompt" )
29-
27+
3028 mock_litellm .assert_called_once_with (
31- "gpt-4" ,
32- [{"role" : "user" , "content" : "Test prompt" }],
33- max_tokens = None
29+ "gpt-4" , [{"role" : "user" , "content" : "Test prompt" }], max_tokens = None
3430 )
35-
31+
3632 @pytest .mark .asyncio
3733 async def test_ping_model_success (self , mock_litellm ):
3834 """Test successful model ping"""
3935 # Create a mock console that can be used in the context
4036 mock_console_instance = MagicMock ()
41-
37+
4238 result = await ping_model ("gpt-4" , mock_console_instance )
43-
39+
4440 # Verify success
4541 assert result is True
46-
42+
4743 # Verify console output
4844 mock_console_instance .print .assert_called_once_with ("[green]✓[/green] gpt-4" )
49-
45+
5046 # Verify LLM was called with validation prompt
5147 mock_litellm .assert_called_once_with (
52- "gpt-4" ,
53- [{"role" : "user" , "content" : VALIDATION_PROMPT }],
54- max_tokens = 1
48+ "gpt-4" , [{"role" : "user" , "content" : VALIDATION_PROMPT }], max_tokens = 1
5549 )
56-
50+
5751 @pytest .mark .asyncio
5852 async def test_ping_model_failure (self , mock_litellm ):
5953 """Test failed model ping"""
6054 # Configure mock to raise exception
6155 mock_litellm .side_effect = Exception ("API Error" )
6256 mock_console_instance = MagicMock ()
63-
57+
6458 result = await ping_model ("invalid-model" , mock_console_instance )
65-
59+
6660 # Verify failure
6761 assert result is False
68-
62+
6963 # Verify error output
7064 mock_console_instance .print .assert_called_once_with (
7165 "[red]✗[/red] invalid-model - API Error"
7266 )
73-
67+
7468 @pytest .mark .asyncio
7569 async def test_bench_model_success (self , mock_litellm , mocker ):
7670 """Test successful benchmark run"""
7771 # Mock time to control duration measurement
78- mock_time = mocker .patch (' tacho.ai.time.time'
72+ mock_time = mocker .patch (" tacho.ai.time.time"
7973 mock_time .side_effect = [100.0 , 102.5 ] # 2.5 second duration
80-
74+
8175 # Configure mock response with usage data (no reasoning tokens)
8276 mock_response = MagicMock ()
8377 mock_usage = MagicMock ()
@@ -86,93 +80,75 @@ async def test_bench_model_success(self, mock_litellm, mocker):
8680 mock_usage .completion_tokens_details = None
8781 mock_response .usage = mock_usage
8882 mock_litellm .return_value = mock_response
89-
83+
9084 duration , tokens = await bench_model ("gpt-4" , 500 )
91-
85+
9286 # Verify results
9387 assert duration == 2.5
9488 assert tokens == 150
95-
89+
9690 # Verify LLM was called correctly
9791 mock_litellm .assert_called_once_with (
98- "gpt-4" ,
99- [{"role" : "user" , "content" : BENCHMARK_PROMPT }],
100- max_tokens = 500
92+ "gpt-4" , [{"role" : "user" , "content" : BENCHMARK_PROMPT }], max_tokens = 500
10193 )
102-
103- @pytest .mark .asyncio
104- async def test_bench_model_no_usage_data (self , mock_litellm , mocker ):
105- """Test benchmark when response has no usage data"""
106- # Mock time
107- mock_time = mocker .patch ('tacho.ai.time.time' )
108- mock_time .side_effect = [100.0 , 101.0 ]
109-
110- # Configure mock response without usage
111- mock_response = MagicMock ()
112- mock_response .usage = None
113- mock_litellm .return_value = mock_response
114-
115- duration , tokens = await bench_model ("gpt-4" , 500 )
116-
117- # Should return 0 tokens when no usage data
118- assert duration == 1.0
119- assert tokens == 0
120-
94+
12195 @pytest .mark .asyncio
12296 async def test_bench_model_exception_handling (self , mock_litellm ):
12397 """Test that exceptions propagate from bench_model"""
12498 mock_litellm .side_effect = Exception ("Network error" )
125-
99+
126100 with pytest .raises (Exception , match = "Network error" ):
127101 await bench_model ("gpt-4" , 500 )
128-
102+
129103 @pytest .mark .asyncio
130104 async def test_bench_model_with_reasoning_tokens (self , mock_litellm , mocker ):
131105 """Test benchmark with reasoning models that have completion_tokens_details"""
132106 # Mock time
133- mock_time = mocker .patch (' tacho.ai.time.time'
107+ mock_time = mocker .patch (" tacho.ai.time.time"
134108 mock_time .side_effect = [100.0 , 103.0 ] # 3 second duration
135-
109+
136110 # Configure mock response with reasoning tokens
137111 mock_response = MagicMock ()
138112 mock_response .usage .completion_tokens = 50 # Regular completion tokens
139-
113+
140114 # Mock completion_tokens_details with reasoning_tokens
141115 mock_details = MagicMock ()
142116 mock_details .reasoning_tokens = 200 # Reasoning tokens
143117 mock_response .usage .completion_tokens_details = mock_details
144-
118+
145119 mock_litellm .return_value = mock_response
146-
120+
147121 duration , tokens = await bench_model ("o1-mini" , 500 )
148-
122+
149123 # Verify results - should include both completion and reasoning tokens
150124 assert duration == 3.0
151125 assert tokens == 250 # 50 completion + 200 reasoning
152-
126+
153127 # Verify LLM was called correctly
154128 mock_litellm .assert_called_once_with (
155- "o1-mini" ,
156- [{"role" : "user" , "content" : BENCHMARK_PROMPT }],
157- max_tokens = 500
129+ "o1-mini" , [{"role" : "user" , "content" : BENCHMARK_PROMPT }], max_tokens = 500
158130 )
159-
131+
160132 @pytest .mark .asyncio
161- async def test_bench_model_with_empty_completion_details (self , mock_litellm , mocker ):
133+ async def test_bench_model_with_empty_completion_details (
134+ self , mock_litellm , mocker
135+ ):
162136 """Test benchmark when completion_tokens_details exists but has no reasoning_tokens"""
163137 # Mock time
164- mock_time = mocker .patch (' tacho.ai.time.time'
138+ mock_time = mocker .patch (" tacho.ai.time.time"
165139 mock_time .side_effect = [100.0 , 102.0 ]
166-
140+
167141 # Configure mock response with completion_tokens_details but no reasoning_tokens
168142 mock_response = MagicMock ()
169143 mock_response .usage .completion_tokens = 100
170- mock_response .usage .completion_tokens_details = MagicMock (spec = []) # No reasoning_tokens attribute
171-
144+ mock_response .usage .completion_tokens_details = MagicMock (
145+ spec = []
146+ ) # No reasoning_tokens attribute
147+
172148 mock_litellm .return_value = mock_response
173-
149+
174150 duration , tokens = await bench_model ("gpt-4" , 500 )
175-
151+
176152 # Should only count regular completion tokens
177153 assert duration == 2.0
178- assert tokens == 100
154+ assert tokens == 100
0 commit comments