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