visualization-agent/test_config_integration.py

"""
Integration tests for image configuration support
Tests that configuration parameters are properly passed to the Gemini API
"""
import os
import pytest
from unittest.mock import Mock, patch, MagicMock
from agent import VisualizationAgent
from google.genai import types


class TestConfigurationIntegration:
    """Test configuration parameter passing to API"""
    
    def setup_method(self):
        """Set up test fixtures"""
        self.api_key = "test-api-key"
        self.agent = VisualizationAgent(api_key=self.api_key)
    
    @patch('agent.genai.Client')
    def test_aspect_ratio_passed_to_api(self, mock_client_class):
        """Test that aspect_ratio is passed to GenerateContentConfig"""
        # Setup mock
        mock_client = Mock()
        mock_client_class.return_value = mock_client
        
        # Create mock response with image data
        mock_part = Mock()
        mock_part.inline_data = Mock()
        mock_image = Mock()
        mock_part.as_image.return_value = mock_image
        
        mock_response = Mock()
        mock_response.parts = [mock_part]
        mock_client.models.generate_content.return_value = mock_response
        
        # Create agent with mocked client
        agent = VisualizationAgent(api_key=self.api_key)
        agent.client = mock_client
        
        # Call with custom aspect ratio
        config = {"aspect_ratio": "4:3"}
        result = agent.generate_image("test prompt", config)
        
        # Verify the API was called with correct config
        assert mock_client.models.generate_content.called
        call_args = mock_client.models.generate_content.call_args
        
        # Check that config parameter was passed
        assert 'config' in call_args.kwargs
        config_arg = call_args.kwargs['config']
        
        # Verify it's a GenerateContentConfig with image_config
        assert isinstance(config_arg, types.GenerateContentConfig)
        assert hasattr(config_arg, 'image_config')
        assert config_arg.image_config.aspect_ratio == "4:3"
    
    @patch('agent.genai.Client')
    def test_image_size_passed_to_api(self, mock_client_class):
        """Test that image_size is passed to GenerateContentConfig"""
        # Setup mock
        mock_client = Mock()
        mock_client_class.return_value = mock_client
        
        # Create mock response with image data
        mock_part = Mock()
        mock_part.inline_data = Mock()
        mock_image = Mock()
        mock_part.as_image.return_value = mock_image
        
        mock_response = Mock()
        mock_response.parts = [mock_part]
        mock_client.models.generate_content.return_value = mock_response
        
        # Create agent with mocked client
        agent = VisualizationAgent(api_key=self.api_key)
        agent.client = mock_client
        
        # Call with custom image size
        config = {"image_size": "2K"}
        result = agent.generate_image("test prompt", config)
        
        # Verify the API was called with correct config
        assert mock_client.models.generate_content.called
        call_args = mock_client.models.generate_content.call_args
        
        # Check that config parameter was passed
        assert 'config' in call_args.kwargs
        config_arg = call_args.kwargs['config']
        
        # Verify it's a GenerateContentConfig with image_config
        assert isinstance(config_arg, types.GenerateContentConfig)
        assert hasattr(config_arg, 'image_config')
        assert config_arg.image_config.image_size == "2K"
    
    @patch('agent.genai.Client')
    def test_model_selection_passed_to_api(self, mock_client_class):
        """Test that model selection is passed to generate_content"""
        # Setup mock
        mock_client = Mock()
        mock_client_class.return_value = mock_client
        
        # Create mock response with image data
        mock_part = Mock()
        mock_part.inline_data = Mock()
        mock_image = Mock()
        mock_part.as_image.return_value = mock_image
        
        mock_response = Mock()
        mock_response.parts = [mock_part]
        mock_client.models.generate_content.return_value = mock_response
        
        # Create agent with mocked client
        agent = VisualizationAgent(api_key=self.api_key)
        agent.client = mock_client
        
        # Call with custom model
        config = {"model": "gemini-3-pro-image-preview"}
        result = agent.generate_image("test prompt", config)
        
        # Verify the API was called with correct model
        assert mock_client.models.generate_content.called
        call_args = mock_client.models.generate_content.call_args
        
        # Check that model parameter was passed
        assert 'model' in call_args.kwargs
        assert call_args.kwargs['model'] == "gemini-3-pro-image-preview"
    
    @patch('agent.genai.Client')
    def test_all_config_parameters_passed(self, mock_client_class):
        """Test that all configuration parameters are passed correctly"""
        # Setup mock
        mock_client = Mock()
        mock_client_class.return_value = mock_client
        
        # Create mock response with image data
        mock_part = Mock()
        mock_part.inline_data = Mock()
        mock_image = Mock()
        mock_part.as_image.return_value = mock_image
        
        mock_response = Mock()
        mock_response.parts = [mock_part]
        mock_client.models.generate_content.return_value = mock_response
        
        # Create agent with mocked client
        agent = VisualizationAgent(api_key=self.api_key)
        agent.client = mock_client
        
        # Call with full config
        config = {
            "model": "gemini-3-pro-image-preview",
            "aspect_ratio": "21:9",
            "image_size": "4K",
            "use_search": False
        }
        result = agent.generate_image("test prompt", config)
        
        # Verify the API was called with all parameters
        assert mock_client.models.generate_content.called
        call_args = mock_client.models.generate_content.call_args
        
        # Check model
        assert call_args.kwargs['model'] == "gemini-3-pro-image-preview"
        
        # Check config
        config_arg = call_args.kwargs['config']
        assert isinstance(config_arg, types.GenerateContentConfig)
        assert config_arg.image_config.aspect_ratio == "21:9"
        assert config_arg.image_config.image_size == "4K"
    
    @patch('agent.genai.Client')
    def test_default_config_when_none_provided(self, mock_client_class):
        """Test that default configuration is used when none provided"""
        # Setup mock
        mock_client = Mock()
        mock_client_class.return_value = mock_client
        
        # Create mock response with image data
        mock_part = Mock()
        mock_part.inline_data = Mock()
        mock_image = Mock()
        mock_part.as_image.return_value = mock_image
        
        mock_response = Mock()
        mock_response.parts = [mock_part]
        mock_client.models.generate_content.return_value = mock_response
        
        # Create agent with mocked client
        agent = VisualizationAgent(api_key=self.api_key)
        agent.client = mock_client
        
        # Call without config
        result = agent.generate_image("test prompt")
        
        # Verify the API was called with defaults
        assert mock_client.models.generate_content.called
        call_args = mock_client.models.generate_content.call_args
        
        # Check default model
        assert call_args.kwargs['model'] == "gemini-2.5-flash-image"
        
        # Check default config
        config_arg = call_args.kwargs['config']
        assert isinstance(config_arg, types.GenerateContentConfig)
        assert config_arg.image_config.aspect_ratio == "16:9"
        assert config_arg.image_config.image_size == "1K"


if __name__ == "__main__":
    pytest.main([__file__, "-v"])