gpt-edit/tests/test_validation.py

"""Unit tests for validation utilities"""

import pytest
from src.utils.validation import (
    validate_edit_parameters,
    validate_batch_parameters,
    sanitize_prompt,
    validate_api_response
)
import tempfile
from PIL import Image


def create_temp_image():
    """Helper to create a temporary image file"""
    with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as f:
        img = Image.new('RGBA', (100, 100), color=(255, 0, 0, 255))
        img.save(f.name, 'PNG')
        return f.name


def test_validate_edit_parameters_valid():
    """Test validation of valid edit parameters"""
    temp_image = create_temp_image()
    
    params = {
        'image_path': temp_image,
        'prompt': 'Make it blue'
    }
    
    is_valid, error_msg = validate_edit_parameters(params)
    assert is_valid is True
    assert error_msg is None


def test_validate_edit_parameters_missing_required():
    """Test validation with missing required fields"""
    params = {
        'prompt': 'Make it blue'
    }
    
    is_valid, error_msg = validate_edit_parameters(params)
    assert is_valid is False
    assert 'image_path' in error_msg


def test_validate_edit_parameters_invalid_image_path():
    """Test validation with non-existent image"""
    params = {
        'image_path': '/nonexistent/image.png',
        'prompt': 'Make it blue'
    }
    
    is_valid, error_msg = validate_edit_parameters(params)
    assert is_valid is False
    assert 'not found' in error_msg


def test_validate_edit_parameters_invalid_background():
    """Test validation with invalid background option"""
    temp_image = create_temp_image()
    
    params = {
        'image_path': temp_image,
        'prompt': 'Make it blue',
        'background': 'invalid_option'
    }
    
    is_valid, error_msg = validate_edit_parameters(params)
    assert is_valid is False
    assert 'background' in error_msg


def test_validate_edit_parameters_invalid_size():
    """Test validation with invalid size"""
    temp_image = create_temp_image()
    
    params = {
        'image_path': temp_image,
        'prompt': 'Make it blue',
        'size': '2048x2048'  # Not supported
    }
    
    is_valid, error_msg = validate_edit_parameters(params)
    assert is_valid is False
    assert 'size' in error_msg


def test_validate_batch_parameters_valid():
    """Test validation of valid batch parameters"""
    temp_image1 = create_temp_image()
    temp_image2 = create_temp_image()
    
    batch_params = [
        {'image_path': temp_image1, 'prompt': 'Make it blue'},
        {'image_path': temp_image2, 'prompt': 'Make it green'}
    ]
    
    is_valid, error_msg = validate_batch_parameters(batch_params)
    assert is_valid is True
    assert error_msg is None


def test_validate_batch_parameters_empty():
    """Test validation with empty batch"""
    batch_params = []
    
    is_valid, error_msg = validate_batch_parameters(batch_params)
    assert is_valid is False
    assert 'No images' in error_msg


def test_validate_batch_parameters_too_many():
    """Test validation with too many images"""
    temp_image = create_temp_image()
    
    # Create 17 items (max is 16)
    batch_params = [
        {'image_path': temp_image, 'prompt': f'Edit {i}'}
        for i in range(17)
    ]
    
    is_valid, error_msg = validate_batch_parameters(batch_params)
    assert is_valid is False
    assert 'Too many' in error_msg


def test_validate_batch_parameters_invalid_item():
    """Test validation with invalid item in batch"""
    temp_image = create_temp_image()
    
    batch_params = [
        {'image_path': temp_image, 'prompt': 'Valid'},
        {'image_path': '/invalid/path.png', 'prompt': 'Invalid'}
    ]
    
    is_valid, error_msg = validate_batch_parameters(batch_params)
    assert is_valid is False
    assert 'Item 2' in error_msg


def test_sanitize_prompt_whitespace():
    """Test prompt sanitization for whitespace"""
    prompt = "  Make   it   blue   "
    sanitized = sanitize_prompt(prompt)
    
    assert sanitized == "Make it blue"


def test_sanitize_prompt_truncation():
    """Test prompt truncation for long prompts"""
    prompt = "x" * 2000  # Very long prompt
    sanitized = sanitize_prompt(prompt)
    
    assert len(sanitized) == 1000


def test_validate_api_response_valid():
    """Test validation of valid API response"""
    # Mock a valid response
    class MockResponse:
        def __init__(self):
            self.data = [MockData()]
    
    class MockData:
        def __init__(self):
            self.b64_json = "base64encodeddata"
    
    response = MockResponse()
    assert validate_api_response(response) is True


def test_validate_api_response_invalid():
    """Test validation of invalid API responses"""
    # None response
    assert validate_api_response(None) is False
    
    # Response without data attribute
    class BadResponse1:
        pass
    assert validate_api_response(BadResponse1()) is False
    
    # Response with empty data
    class BadResponse2:
        def __init__(self):
            self.data = []
    assert validate_api_response(BadResponse2()) is False
    
    # Response without b64_json
    class BadResponse3:
        def __init__(self):
            self.data = [object()]
    assert validate_api_response(BadResponse3()) is False


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