visualization-agent/test_error_handling.py

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