File size: 8,002 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 |
"""
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"])
|