visualization-agent/test_input_validation.py

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