File size: 14,316 Bytes
9b24c4d |
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 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 |
"""
Tests for comprehensive error handling in VisualizationAgent
"""
import os
import pytest
from unittest.mock import Mock, patch, MagicMock
from PIL import Image
from agent import VisualizationAgent
class TestNetworkErrorHandling:
"""Test network error handling with retry_possible=true"""
def test_connection_error_returns_network_error(self):
"""Test that connection errors are properly handled"""
agent = VisualizationAgent(api_key="test-key")
# Mock the client to raise a connection error
with patch.object(agent.client.models, 'generate_content') as mock_generate:
mock_generate.side_effect = Exception("Connection refused")
result = agent.generate_image("test prompt")
assert result["success"] is False
assert result["error"]["code"] == "NETWORK_ERROR"
assert result["error"]["retry_possible"] is True
assert "network" in result["error"]["message"].lower()
def test_timeout_error_returns_network_error(self):
"""Test that timeout errors are properly handled"""
agent = VisualizationAgent(api_key="test-key")
with patch.object(agent.client.models, 'generate_content') as mock_generate:
mock_generate.side_effect = Exception("Request timed out after 30 seconds")
result = agent.generate_image("test prompt")
assert result["success"] is False
assert result["error"]["code"] == "NETWORK_ERROR"
assert result["error"]["retry_possible"] is True
def test_unreachable_error_returns_network_error(self):
"""Test that unreachable errors are properly handled"""
agent = VisualizationAgent(api_key="test-key")
with patch.object(agent.client.models, 'generate_content') as mock_generate:
mock_generate.side_effect = Exception("Network unreachable")
result = agent.generate_image("test prompt")
assert result["success"] is False
assert result["error"]["code"] == "NETWORK_ERROR"
assert result["error"]["retry_possible"] is True
class TestAuthenticationErrorHandling:
"""Test authentication error handling with descriptive messages"""
def test_invalid_api_key_returns_auth_error(self):
"""Test that invalid API key errors are properly handled"""
agent = VisualizationAgent(api_key="test-key")
with patch.object(agent.client.models, 'generate_content') as mock_generate:
mock_generate.side_effect = Exception("Invalid API key provided")
result = agent.generate_image("test prompt")
assert result["success"] is False
assert result["error"]["code"] == "AUTH_ERROR"
assert result["error"]["retry_possible"] is False
assert "API key" in result["error"]["message"] or "credential" in result["error"]["message"].lower()
def test_unauthorized_error_returns_auth_error(self):
"""Test that 401 unauthorized errors are properly handled"""
agent = VisualizationAgent(api_key="test-key")
with patch.object(agent.client.models, 'generate_content') as mock_generate:
mock_generate.side_effect = Exception("401 Unauthorized")
result = agent.generate_image("test prompt")
assert result["success"] is False
assert result["error"]["code"] == "AUTH_ERROR"
assert result["error"]["retry_possible"] is False
def test_authentication_failed_returns_auth_error(self):
"""Test that authentication failures are properly handled"""
agent = VisualizationAgent(api_key="test-key")
with patch.object(agent.client.models, 'generate_content') as mock_generate:
mock_generate.side_effect = Exception("Authentication failed")
result = agent.generate_image("test prompt")
assert result["success"] is False
assert result["error"]["code"] == "AUTH_ERROR"
assert result["error"]["retry_possible"] is False
class TestRateLimitErrorHandling:
"""Test rate limit error handling with retry information"""
def test_rate_limit_error_returns_rate_limit(self):
"""Test that rate limit errors are properly handled"""
agent = VisualizationAgent(api_key="test-key")
with patch.object(agent.client.models, 'generate_content') as mock_generate:
mock_generate.side_effect = Exception("Rate limit exceeded")
result = agent.generate_image("test prompt")
assert result["success"] is False
assert result["error"]["code"] == "RATE_LIMIT"
assert result["error"]["retry_possible"] is True
assert "details" in result["error"]
assert "retry_after" in result["error"]["details"]
def test_quota_exceeded_returns_rate_limit(self):
"""Test that quota exceeded errors are properly handled"""
agent = VisualizationAgent(api_key="test-key")
with patch.object(agent.client.models, 'generate_content') as mock_generate:
mock_generate.side_effect = Exception("Quota exceeded for this API")
result = agent.generate_image("test prompt")
assert result["success"] is False
assert result["error"]["code"] == "RATE_LIMIT"
assert result["error"]["retry_possible"] is True
def test_429_error_returns_rate_limit(self):
"""Test that 429 errors are properly handled"""
agent = VisualizationAgent(api_key="test-key")
with patch.object(agent.client.models, 'generate_content') as mock_generate:
mock_generate.side_effect = Exception("429 Too Many Requests")
result = agent.generate_image("test prompt")
assert result["success"] is False
assert result["error"]["code"] == "RATE_LIMIT"
assert result["error"]["retry_possible"] is True
# Check that retry information is present in message or details
message_lower = result["error"]["message"].lower()
assert "retry" in message_lower or "wait" in message_lower
class TestMissingImageDataHandling:
"""Test handling of missing image data in response"""
def test_no_image_data_in_response(self):
"""Test that missing image data is properly detected"""
agent = VisualizationAgent(api_key="test-key")
# Mock response with no image data
mock_response = Mock()
mock_response.parts = []
with patch.object(agent.client.models, 'generate_content') as mock_generate:
mock_generate.return_value = mock_response
result = agent.generate_image("test prompt")
assert result["success"] is False
assert result["error"]["code"] == "NO_IMAGE_DATA"
assert result["error"]["retry_possible"] is True
assert "no image data" in result["error"]["message"].lower()
def test_response_with_no_inline_data(self):
"""Test that response without inline_data is properly handled"""
agent = VisualizationAgent(api_key="test-key")
# Mock response with parts but no inline_data
mock_part = Mock()
mock_part.inline_data = None
mock_response = Mock()
mock_response.parts = [mock_part]
with patch.object(agent.client.models, 'generate_content') as mock_generate:
mock_generate.return_value = mock_response
result = agent.generate_image("test prompt")
assert result["success"] is False
assert result["error"]["code"] == "NO_IMAGE_DATA"
assert result["error"]["retry_possible"] is True
class TestFileSystemErrorHandling:
"""Test file system error handling during save"""
def test_permission_denied_error(self):
"""Test that permission denied errors are handled gracefully (stateless mode)"""
# Mock os.makedirs to avoid creating directory during initialization
with patch('os.makedirs'):
agent = VisualizationAgent(api_key="test-key", output_dir="/root/no-permission")
# Mock successful API call
mock_part = Mock()
mock_part.inline_data = True
mock_image = Image.new('RGB', (100, 100), color='red')
mock_part.as_image.return_value = mock_image
mock_response = Mock()
mock_response.parts = [mock_part]
with patch.object(agent.client.models, 'generate_content') as mock_generate:
mock_generate.return_value = mock_response
# Mock os.access to return False (no write permission)
with patch('os.access', return_value=False):
result = agent.generate_image("test prompt")
# In stateless mode, file save failure is logged but doesn't fail the request
# The base64 encoding is the primary delivery method
assert result["success"] is True
assert result["image_base64"] is not None
# image_path may be None since save failed
assert result["image_path"] is None
def test_disk_full_error(self):
"""Test that disk full errors are properly handled"""
agent = VisualizationAgent(api_key="test-key")
# Mock successful API call
mock_part = Mock()
mock_part.inline_data = True
mock_image = Mock()
mock_image.save.side_effect = OSError("No space left on device")
mock_part.as_image.return_value = mock_image
mock_response = Mock()
mock_response.parts = [mock_part]
with patch.object(agent.client.models, 'generate_content') as mock_generate:
mock_generate.return_value = mock_response
result = agent.generate_image("test prompt")
assert result["success"] is False
assert result["error"]["code"] == "FILE_SYSTEM_ERROR"
assert result["error"]["retry_possible"] is False
assert "disk" in result["error"]["message"].lower() or "space" in result["error"]["message"].lower()
def test_read_only_filesystem_error(self):
"""Test that read-only filesystem errors are properly handled"""
agent = VisualizationAgent(api_key="test-key")
# Mock successful API call
mock_part = Mock()
mock_part.inline_data = True
mock_image = Mock()
mock_image.save.side_effect = OSError("Read-only file system")
mock_part.as_image.return_value = mock_image
mock_response = Mock()
mock_response.parts = [mock_part]
with patch.object(agent.client.models, 'generate_content') as mock_generate:
mock_generate.return_value = mock_response
result = agent.generate_image("test prompt")
assert result["success"] is False
assert result["error"]["code"] == "FILE_SYSTEM_ERROR"
assert result["error"]["retry_possible"] is False
assert "read-only" in result["error"]["message"].lower()
class TestStandardResponseFormat:
"""Test that all errors follow standard response format"""
def test_error_response_has_required_fields(self):
"""Test that error responses have all required fields"""
agent = VisualizationAgent(api_key="test-key")
with patch.object(agent.client.models, 'generate_content') as mock_generate:
mock_generate.side_effect = Exception("Test error")
result = agent.generate_image("test prompt")
# Check standard response format
assert "success" in result
assert result["success"] is False
assert "error" in result
assert "code" in result["error"]
assert "message" in result["error"]
assert "retry_possible" in result["error"]
def test_success_response_has_required_fields(self):
"""Test that success responses have all required fields"""
agent = VisualizationAgent(api_key="test-key")
# Mock successful generation
mock_part = Mock()
mock_part.inline_data = True
mock_image = Mock()
mock_part.as_image.return_value = mock_image
mock_response = Mock()
mock_response.parts = [mock_part]
with patch.object(agent.client.models, 'generate_content') as mock_generate:
mock_generate.return_value = mock_response
with patch.object(agent, '_save_image', return_value="/path/to/image.png"):
result = agent.generate_image("test prompt")
# Check standard response format
assert "success" in result
assert result["success"] is True
assert "image_path" in result
assert result["image_path"] == "/path/to/image.png"
class TestContentPolicyViolations:
"""Test content policy violation handling"""
def test_content_policy_violation_error(self):
"""Test that content policy violations are properly handled"""
agent = VisualizationAgent(api_key="test-key")
with patch.object(agent.client.models, 'generate_content') as mock_generate:
mock_generate.side_effect = Exception("Content policy violation detected")
result = agent.generate_image("inappropriate prompt")
assert result["success"] is False
assert result["error"]["code"] == "CONTENT_POLICY_VIOLATION"
assert result["error"]["retry_possible"] is False
assert "content policy" in result["error"]["message"].lower()
if __name__ == "__main__":
pytest.main([__file__, "-v"])
|