tests/test_image_processor.py

"""Tests for image_processor module."""

import pytest
import json
import zipfile
from pathlib import Path
from PIL import Image
from unittest.mock import patch, mock_open

from src.image_processor import (
    save_image,
    create_thumbnail,
    validate_image_format,
    create_zip,
    optimize_image_for_web,
    load_metadata,
)


class TestImageProcessor:
    """Test cases for image processing functions."""
    
    def test_create_thumbnail(self):
        """Test thumbnail creation."""
        # Create a test image
        test_image = Image.new("RGB", (1024, 1024), color="red")
        
        thumbnail = create_thumbnail(test_image, (256, 256))
        
        assert thumbnail.size[0] <= 256
        assert thumbnail.size[1] <= 256
        assert isinstance(thumbnail, Image.Image)
    
    def test_validate_image_format(self, tmp_path):
        """Test image format validation."""
        # Test valid formats
        valid_paths = [
            tmp_path / "test.png",
            tmp_path / "test.jpg",
            tmp_path / "test.jpeg",
            tmp_path / "test.webp",
        ]
        
        for path in valid_paths:
            path.touch()
            assert validate_image_format(path) is True
        
        # Test invalid format
        invalid_path = tmp_path / "test.txt"
        invalid_path.touch()
        assert validate_image_format(invalid_path) is False
    
    def test_optimize_image_for_web(self):
        """Test image optimization for web."""
        # Create large image
        large_image = Image.new("RGB", (4096, 4096), color="blue")
        
        optimized = optimize_image_for_web(large_image, max_size=2048)
        
        assert optimized.size[0] <= 2048
        assert optimized.size[1] <= 2048
        
        # Test with small image (should not change)
        small_image = Image.new("RGB", (512, 512), color="green")
        optimized_small = optimize_image_for_web(small_image, max_size=2048)
        
        assert optimized_small.size == small_image.size
    
    def test_save_image(self, tmp_path, monkeypatch):
        """Test saving image with metadata."""
        # Set output directory
        monkeypatch.setattr("src.image_processor.OUTPUT_DIR", tmp_path)
        
        # Create test image
        test_image = Image.new("RGB", (256, 256), color="red")
        
        params = {
            "guidance_scale": 7.5,
            "num_inference_steps": 50,
            "seed": 12345,
        }
        
        image_path = save_image(
            test_image,
            prompt="test prompt",
            params=params,
        )
        
        # Check image was saved
        assert image_path.exists()
        assert image_path.suffix == ".png"
        
        # Check metadata was saved
        metadata_path = image_path.with_suffix(".json")
        assert metadata_path.exists()
        
        with open(metadata_path, "r") as f:
            metadata = json.load(f)
        
        assert metadata["prompt"] == "test prompt"
        assert metadata["params"] == params
    
    def test_create_zip(self, tmp_path):
        """Test ZIP creation."""
        # Create test files
        test_files = []
        for i in range(3):
            test_file = tmp_path / f"test_{i}.txt"
            test_file.write_text(f"Content {i}")
            test_files.append(test_file)
        
        # Create ZIP
        zip_path = tmp_path / "test.zip"
        create_zip(test_files, zip_path)
        
        # Verify ZIP exists
        assert zip_path.exists()
        
        # Verify contents
        with zipfile.ZipFile(zip_path, "r") as zipf:
            names = zipf.namelist()
            assert len(names) == 3
            assert all(f"test_{i}.txt" in names for i in range(3))
    
    def test_load_metadata(self, tmp_path):
        """Test loading metadata."""
        # Create test metadata
        metadata_path = tmp_path / "test.json"
        test_metadata = {
            "timestamp": 1234567890,
            "prompt": "test prompt",
            "params": {"seed": 12345},
        }
        
        with open(metadata_path, "w") as f:
            json.dump(test_metadata, f)
        
        # Test loading
        loaded = load_metadata(metadata_path)
        
        assert loaded is not None
        assert loaded["prompt"] == "test prompt"
        assert loaded["params"]["seed"] == 12345
        
        # Test with non-existent file
        non_existent = tmp_path / "nonexistent.json"
        assert load_metadata(non_existent) is None


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