""" 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"])