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