1+ """
2+ Tests for the Gemini Model error handling scenarios.
3+ """
4+ import json
5+ from unittest .mock import patch , MagicMock
6+
7+ import pytest
8+ from rich .console import Console
9+
10+ from cli_code .models .gemini import GeminiModel
11+ from cli_code .tools import AVAILABLE_TOOLS
12+
13+
14+ class TestGeminiModelErrorHandling :
15+ """Tests for error handling in GeminiModel."""
16+
17+ @pytest .fixture
18+ def mock_generative_model (self ):
19+ """Mock the Gemini generative model."""
20+ with patch ("cli_code.models.gemini.generative_models.GenerativeModel" ) as mock_model :
21+ mock_instance = MagicMock ()
22+ mock_model .return_value = mock_instance
23+ yield mock_instance
24+
25+ @pytest .fixture
26+ def gemini_model (self , mock_generative_model ):
27+ """Create a GeminiModel instance with mocked dependencies."""
28+ console = Console ()
29+ with patch ("cli_code.models.gemini.generative_models" ) as mock_gm :
30+ # Configure the mock
31+ mock_gm .GenerativeModel = MagicMock ()
32+ mock_gm .GenerativeModel .return_value = mock_generative_model
33+
34+ # Create the model
35+ model = GeminiModel (api_key = "fake_api_key" , console = console , model_name = "gemini-pro" )
36+ yield model
37+
38+ @patch ("cli_code.models.gemini.generative_models" )
39+ def test_initialization_error (self , mock_gm ):
40+ """Test error handling during initialization."""
41+ # Make the GenerativeModel constructor raise an exception
42+ mock_gm .GenerativeModel .side_effect = Exception ("API initialization error" )
43+
44+ # Create a console for the model
45+ console = Console ()
46+
47+ # Attempt to create the model - should raise an error
48+ with pytest .raises (Exception ) as excinfo :
49+ GeminiModel (api_key = "fake_api_key" , console = console , model_name = "gemini-pro" )
50+
51+ # Verify the error message
52+ assert "API initialization error" in str (excinfo .value )
53+
54+ def test_empty_prompt_error (self , gemini_model , mock_generative_model ):
55+ """Test error handling when an empty prompt is provided."""
56+ # Call generate with an empty prompt
57+ result = gemini_model .generate ("" )
58+
59+ # Verify error message is returned
60+ assert result is not None
61+ assert "empty prompt" in result .lower ()
62+
63+ # Verify that no API call was made
64+ mock_generative_model .generate_content .assert_not_called ()
65+
66+ def test_api_error_handling (self , gemini_model , mock_generative_model ):
67+ """Test handling of API errors during generation."""
68+ # Make the API call raise an exception
69+ mock_generative_model .generate_content .side_effect = Exception ("API error" )
70+
71+ # Call generate
72+ result = gemini_model .generate ("Test prompt" )
73+
74+ # Verify error message is returned
75+ assert result is not None
76+ assert "error" in result .lower ()
77+ assert "api error" in result .lower ()
78+
79+ def test_rate_limit_error_handling (self , gemini_model , mock_generative_model ):
80+ """Test handling of rate limit errors."""
81+ # Create a rate limit error
82+ rate_limit_error = Exception ("Rate limit exceeded" )
83+ mock_generative_model .generate_content .side_effect = rate_limit_error
84+
85+ # Call generate
86+ result = gemini_model .generate ("Test prompt" )
87+
88+ # Verify rate limit error message is returned
89+ assert result is not None
90+ assert "rate limit" in result .lower () or "quota" in result .lower ()
91+
92+ def test_invalid_api_key_error (self , gemini_model , mock_generative_model ):
93+ """Test handling of invalid API key errors."""
94+ # Create an authentication error
95+ auth_error = Exception ("Invalid API key" )
96+ mock_generative_model .generate_content .side_effect = auth_error
97+
98+ # Call generate
99+ result = gemini_model .generate ("Test prompt" )
100+
101+ # Verify authentication error message is returned
102+ assert result is not None
103+ assert "api key" in result .lower () or "authentication" in result .lower ()
104+
105+ def test_model_not_found_error (self , mock_generative_model ):
106+ """Test handling of model not found errors."""
107+ # Create a console for the model
108+ console = Console ()
109+
110+ # Create the model with an invalid model name
111+ with patch ("cli_code.models.gemini.generative_models" ) as mock_gm :
112+ mock_gm .GenerativeModel .side_effect = Exception ("Model not found: nonexistent-model" )
113+
114+ # Attempt to create the model
115+ with pytest .raises (Exception ) as excinfo :
116+ GeminiModel (api_key = "fake_api_key" , console = console , model_name = "nonexistent-model" )
117+
118+ # Verify the error message
119+ assert "model not found" in str (excinfo .value ).lower ()
120+
121+ @patch ("cli_code.models.gemini.get_tool" )
122+ def test_tool_execution_error (self , mock_get_tool , gemini_model , mock_generative_model ):
123+ """Test handling of errors during tool execution."""
124+ # Configure the mock to return a response with a function call
125+ mock_response = MagicMock ()
126+ mock_parts = [MagicMock ()]
127+ mock_parts [0 ].text = None # No text
128+ mock_parts [0 ].function_call = MagicMock ()
129+ mock_parts [0 ].function_call .name = "test_tool"
130+ mock_parts [0 ].function_call .args = {"arg1" : "value1" }
131+
132+ mock_response .candidates = [MagicMock ()]
133+ mock_response .candidates [0 ].content .parts = mock_parts
134+
135+ mock_generative_model .generate_content .return_value = mock_response
136+
137+ # Make the tool execution raise an error
138+ mock_tool = MagicMock ()
139+ mock_tool .execute .side_effect = Exception ("Tool execution error" )
140+ mock_get_tool .return_value = mock_tool
141+
142+ # Call generate
143+ result = gemini_model .generate ("Use the test_tool" )
144+
145+ # Verify tool error is handled and included in the response
146+ assert result is not None
147+ assert "error" in result .lower ()
148+ assert "tool execution error" in result .lower ()
149+
150+ def test_invalid_function_call_format (self , gemini_model , mock_generative_model ):
151+ """Test handling of invalid function call format."""
152+ # Configure the mock to return a response with an invalid function call
153+ mock_response = MagicMock ()
154+ mock_parts = [MagicMock ()]
155+ mock_parts [0 ].text = None # No text
156+ mock_parts [0 ].function_call = MagicMock ()
157+ mock_parts [0 ].function_call .name = "nonexistent_tool" # Tool doesn't exist
158+ mock_parts [0 ].function_call .args = {"arg1" : "value1" }
159+
160+ mock_response .candidates = [MagicMock ()]
161+ mock_response .candidates [0 ].content .parts = mock_parts
162+
163+ mock_generative_model .generate_content .return_value = mock_response
164+
165+ # Call generate
166+ result = gemini_model .generate ("Use a tool" )
167+
168+ # Verify invalid tool error is handled
169+ assert result is not None
170+ assert "tool not found" in result .lower () or "nonexistent_tool" in result .lower ()
171+
172+ def test_missing_required_args (self , gemini_model , mock_generative_model ):
173+ """Test handling of function calls with missing required arguments."""
174+ # First mock getting a real tool from AVAILABLE_TOOLS
175+ test_tool = None
176+ for tool in AVAILABLE_TOOLS :
177+ if tool .required_args : # Find a tool with required args
178+ test_tool = tool
179+ break
180+
181+ if not test_tool :
182+ pytest .skip ("No tools with required arguments found for testing" )
183+
184+ # Configure the mock to return a response with a function call missing required args
185+ mock_response = MagicMock ()
186+ mock_parts = [MagicMock ()]
187+ mock_parts [0 ].text = None # No text
188+ mock_parts [0 ].function_call = MagicMock ()
189+ mock_parts [0 ].function_call .name = test_tool .name
190+ mock_parts [0 ].function_call .args = {} # Empty args, missing required ones
191+
192+ mock_response .candidates = [MagicMock ()]
193+ mock_response .candidates [0 ].content .parts = mock_parts
194+
195+ mock_generative_model .generate_content .return_value = mock_response
196+
197+ # Patch the get_tool function to return our test tool
198+ with patch ("cli_code.models.gemini.get_tool" ) as mock_get_tool :
199+ mock_get_tool .return_value = test_tool
200+
201+ # Call generate
202+ result = gemini_model .generate ("Use a tool" )
203+
204+ # Verify missing args error is handled
205+ assert result is not None
206+ assert "missing" in result .lower () or "required" in result .lower () or "argument" in result .lower ()
207+
208+ def test_handling_empty_response (self , gemini_model , mock_generative_model ):
209+ """Test handling of empty response from the API."""
210+ # Configure the mock to return an empty response
211+ mock_response = MagicMock ()
212+ mock_response .candidates = [] # No candidates
213+
214+ mock_generative_model .generate_content .return_value = mock_response
215+
216+ # Call generate
217+ result = gemini_model .generate ("Test prompt" )
218+
219+ # Verify empty response is handled
220+ assert result is not None
221+ assert "empty response" in result .lower () or "no response" in result .lower ()
0 commit comments