Spaces:
Running on Zero
Running on Zero
File size: 4,610 Bytes
f0d9a3e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 | """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__])
|