""" Test input validation for VisualizationAgent Tests Requirements 1.5: Input validation """ import os import sys import pytest from pathlib import Path # Add parent directory to path for imports sys.path.append(str(Path(__file__).parent.parent)) from agent import VisualizationAgent class TestPromptValidation: """Test prompt validation""" def setup_method(self): """Set up test agent""" # Use a dummy API key for validation tests (won't make actual API calls) self.agent = VisualizationAgent( api_key="test_key_for_validation", output_dir="./test_output" ) def test_empty_prompt_returns_error(self): """Test that empty prompt returns error""" result = self.agent.generate_image("") assert result["success"] is False assert result["error"]["code"] == "INVALID_INPUT" assert "empty" in result["error"]["message"].lower() assert result["error"]["retry_possible"] is False def test_whitespace_only_prompt_returns_error(self): """Test that whitespace-only prompt returns error""" result = self.agent.generate_image(" \n\t ") assert result["success"] is False assert result["error"]["code"] == "INVALID_INPUT" assert "empty" in result["error"]["message"].lower() def test_none_prompt_returns_error(self): """Test that None prompt returns error""" result = self.agent.generate_image(None) assert result["success"] is False assert result["error"]["code"] == "INVALID_INPUT" assert "none" in result["error"]["message"].lower() def test_very_long_prompt_returns_error(self): """Test that excessively long prompt returns error""" long_prompt = "a" * 5001 result = self.agent.generate_image(long_prompt) assert result["success"] is False assert result["error"]["code"] == "INVALID_INPUT" assert "maximum length" in result["error"]["message"].lower() class TestConfigValidation: """Test configuration parameter validation""" def setup_method(self): """Set up test agent""" self.agent = VisualizationAgent( api_key="test_key_for_validation", output_dir="./test_output" ) def test_invalid_model_returns_error(self): """Test that invalid model name returns error""" config = {"model": "invalid-model"} result = self.agent.generate_image("test prompt", config) assert result["success"] is False assert result["error"]["code"] == "INVALID_CONFIG" assert "invalid model" in result["error"]["message"].lower() def test_valid_models_accepted(self): """Test that valid model names are accepted (validation passes)""" valid_models = ["gemini-2.5-flash-image", "gemini-3-pro-image-preview"] for model in valid_models: config = {"model": model} # Validation should pass (will fail later on API call with dummy key) result = self.agent._validate_config(config) assert result is None, f"Model {model} should be valid" def test_invalid_aspect_ratio_returns_error(self): """Test that invalid aspect ratio returns error""" config = {"aspect_ratio": "99:1"} result = self.agent.generate_image("test prompt", config) assert result["success"] is False assert result["error"]["code"] == "INVALID_CONFIG" assert "invalid aspect_ratio" in result["error"]["message"].lower() def test_valid_aspect_ratios_accepted(self): """Test that valid aspect ratios are accepted""" valid_ratios = ["1:1", "16:9", "4:3", "9:16", "21:9"] for ratio in valid_ratios: config = {"aspect_ratio": ratio} result = self.agent._validate_config(config) assert result is None, f"Aspect ratio {ratio} should be valid" def test_invalid_image_size_returns_error(self): """Test that invalid image size returns error""" config = {"image_size": "8K"} result = self.agent.generate_image("test prompt", config) assert result["success"] is False assert result["error"]["code"] == "INVALID_CONFIG" assert "invalid image_size" in result["error"]["message"].lower() def test_valid_image_sizes_accepted(self): """Test that valid image sizes are accepted""" valid_sizes = ["1K", "2K", "4K"] for size in valid_sizes: # For 4K, need to use the pro model if size == "4K": config = {"image_size": size, "model": "gemini-3-pro-image-preview"} else: config = {"image_size": size} result = self.agent._validate_config(config) assert result is None, f"Image size {size} should be valid" def test_4k_with_flash_model_returns_error(self): """Test that 4K with flash model returns error""" config = { "image_size": "4K", "model": "gemini-2.5-flash-image" } result = self.agent.generate_image("test prompt", config) assert result["success"] is False assert result["error"]["code"] == "INVALID_CONFIG" assert "4k" in result["error"]["message"].lower() assert "gemini-3-pro-image-preview" in result["error"]["message"].lower() def test_invalid_use_search_type_returns_error(self): """Test that non-boolean use_search returns error""" config = {"use_search": "true"} # String instead of boolean result = self.agent.generate_image("test prompt", config) assert result["success"] is False assert result["error"]["code"] == "INVALID_CONFIG" assert "use_search" in result["error"]["message"].lower() assert "boolean" in result["error"]["message"].lower() def test_valid_use_search_accepted(self): """Test that boolean use_search values are accepted""" for value in [True, False]: config = {"use_search": value} result = self.agent._validate_config(config) assert result is None, f"use_search={value} should be valid" def test_invalid_output_format_returns_error(self): """Test that invalid output format returns error""" config = {"output_format": "gif"} result = self.agent.generate_image("test prompt", config) assert result["success"] is False assert result["error"]["code"] == "INVALID_CONFIG" assert "invalid output_format" in result["error"]["message"].lower() def test_valid_output_formats_accepted(self): """Test that valid output formats are accepted""" valid_formats = ["png", "PNG", "jpg", "jpeg", "JPEG"] for fmt in valid_formats: config = {"output_format": fmt} result = self.agent._validate_config(config) assert result is None, f"Output format {fmt} should be valid" def test_non_string_model_returns_error(self): """Test that non-string model returns error""" config = {"model": 123} result = self.agent.generate_image("test prompt", config) assert result["success"] is False assert result["error"]["code"] == "INVALID_CONFIG" assert "model must be a string" in result["error"]["message"].lower() def test_non_string_aspect_ratio_returns_error(self): """Test that non-string aspect_ratio returns error""" config = {"aspect_ratio": 16.9} result = self.agent.generate_image("test prompt", config) assert result["success"] is False assert result["error"]["code"] == "INVALID_CONFIG" assert "aspect ratio must be a string" in result["error"]["message"].lower() def test_non_string_image_size_returns_error(self): """Test that non-string image_size returns error""" config = {"image_size": 1024} result = self.agent.generate_image("test prompt", config) assert result["success"] is False assert result["error"]["code"] == "INVALID_CONFIG" assert "image size must be a string" in result["error"]["message"].lower() def test_multiple_valid_config_params(self): """Test that multiple valid config parameters work together""" config = { "model": "gemini-2.5-flash-image", "aspect_ratio": "16:9", "image_size": "2K", "use_search": False, "output_format": "png" } result = self.agent._validate_config(config) assert result is None, "All valid config parameters should pass validation" if __name__ == "__main__": pytest.main([__file__, "-v"])