229 lines
6.7 KiB
Python
229 lines
6.7 KiB
Python
"""Unit tests for image utilities"""
|
|
|
|
import pytest
|
|
import tempfile
|
|
import os
|
|
from pathlib import Path
|
|
from PIL import Image
|
|
import io
|
|
|
|
from src.utils.image_utils import (
|
|
validate_image_file,
|
|
convert_to_png,
|
|
get_image_dimensions,
|
|
get_image_dimensions_from_bytes,
|
|
ensure_transparent_background,
|
|
ensure_opaque_background,
|
|
save_image,
|
|
create_mask_from_alpha,
|
|
encode_image_base64,
|
|
decode_image_base64
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def temp_image_file():
|
|
"""Create a temporary test image"""
|
|
with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as f:
|
|
# Create a simple test image
|
|
img = Image.new('RGBA', (100, 100), color=(255, 0, 0, 128))
|
|
img.save(f.name, 'PNG')
|
|
yield f.name
|
|
# Cleanup
|
|
if os.path.exists(f.name):
|
|
os.unlink(f.name)
|
|
|
|
|
|
@pytest.fixture
|
|
def large_image_file():
|
|
"""Create a large test image"""
|
|
with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as f:
|
|
# Create a large test image (5000x5000 would be too large)
|
|
img = Image.new('RGBA', (5000, 5000), color=(0, 255, 0, 255))
|
|
img.save(f.name, 'PNG')
|
|
yield f.name
|
|
# Cleanup
|
|
if os.path.exists(f.name):
|
|
os.unlink(f.name)
|
|
|
|
|
|
def test_validate_image_file_valid(temp_image_file):
|
|
"""Test validation of a valid image file"""
|
|
is_valid, size_mb, error_msg = validate_image_file(temp_image_file)
|
|
|
|
assert is_valid is True
|
|
assert size_mb > 0
|
|
assert error_msg is None
|
|
|
|
|
|
def test_validate_image_file_not_exists():
|
|
"""Test validation of non-existent file"""
|
|
is_valid, size_mb, error_msg = validate_image_file("nonexistent.png")
|
|
|
|
assert is_valid is False
|
|
assert size_mb == 0
|
|
assert "not found" in error_msg.lower()
|
|
|
|
|
|
def test_validate_image_file_too_large(large_image_file):
|
|
"""Test validation of oversized image"""
|
|
is_valid, size_mb, error_msg = validate_image_file(large_image_file, max_size_mb=0.001)
|
|
|
|
assert is_valid is False
|
|
assert size_mb > 0.001
|
|
assert "too large" in error_msg.lower()
|
|
|
|
|
|
def test_validate_image_dimensions_too_large():
|
|
"""Test validation of image with dimensions exceeding limits"""
|
|
with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as f:
|
|
img = Image.new('RGBA', (5000, 5000), color=(255, 255, 255, 255))
|
|
img.save(f.name, 'PNG')
|
|
|
|
try:
|
|
is_valid, size_mb, error_msg = validate_image_file(f.name)
|
|
assert is_valid is False
|
|
assert "too large" in error_msg.lower()
|
|
assert "5000x5000" in error_msg
|
|
finally:
|
|
os.unlink(f.name)
|
|
|
|
|
|
def test_convert_to_png_from_jpeg():
|
|
"""Test converting JPEG to PNG"""
|
|
with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as f:
|
|
# Create a JPEG image
|
|
img = Image.new('RGB', (50, 50), color=(255, 0, 0))
|
|
img.save(f.name, 'JPEG')
|
|
|
|
try:
|
|
png_data = convert_to_png(f.name)
|
|
|
|
# Verify it's valid PNG data
|
|
img_converted = Image.open(io.BytesIO(png_data))
|
|
assert img_converted.format == 'PNG'
|
|
assert img_converted.mode == 'RGBA'
|
|
assert img_converted.size == (50, 50)
|
|
finally:
|
|
os.unlink(f.name)
|
|
|
|
|
|
def test_get_image_dimensions(temp_image_file):
|
|
"""Test getting image dimensions"""
|
|
width, height = get_image_dimensions(temp_image_file)
|
|
|
|
assert width == 100
|
|
assert height == 100
|
|
|
|
|
|
def test_get_image_dimensions_from_bytes():
|
|
"""Test getting dimensions from image bytes"""
|
|
img = Image.new('RGBA', (200, 150), color=(0, 0, 255, 255))
|
|
buffer = io.BytesIO()
|
|
img.save(buffer, format='PNG')
|
|
image_data = buffer.getvalue()
|
|
|
|
width, height = get_image_dimensions_from_bytes(image_data)
|
|
|
|
assert width == 200
|
|
assert height == 150
|
|
|
|
|
|
def test_ensure_transparent_background():
|
|
"""Test ensuring transparent background"""
|
|
# Create image with opaque background
|
|
img = Image.new('RGB', (50, 50), color=(255, 255, 255))
|
|
buffer = io.BytesIO()
|
|
img.save(buffer, format='PNG')
|
|
image_data = buffer.getvalue()
|
|
|
|
# Process
|
|
result_data = ensure_transparent_background(image_data)
|
|
|
|
# Verify result has alpha channel
|
|
result_img = Image.open(io.BytesIO(result_data))
|
|
assert result_img.mode == 'RGBA'
|
|
|
|
|
|
def test_ensure_opaque_background():
|
|
"""Test ensuring opaque background"""
|
|
# Create image with transparent areas
|
|
img = Image.new('RGBA', (50, 50), color=(255, 0, 0, 128))
|
|
buffer = io.BytesIO()
|
|
img.save(buffer, format='PNG')
|
|
image_data = buffer.getvalue()
|
|
|
|
# Process
|
|
result_data = ensure_opaque_background(image_data)
|
|
|
|
# Verify result is opaque
|
|
result_img = Image.open(io.BytesIO(result_data))
|
|
assert result_img.mode == 'RGB'
|
|
|
|
|
|
def test_save_image():
|
|
"""Test saving image to file"""
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
# Create test image data
|
|
img = Image.new('RGBA', (50, 50), color=(0, 255, 0, 255))
|
|
buffer = io.BytesIO()
|
|
img.save(buffer, format='PNG')
|
|
image_data = buffer.getvalue()
|
|
|
|
# Save
|
|
output_path = os.path.join(tmpdir, 'test_output.png')
|
|
success = save_image(image_data, output_path)
|
|
|
|
assert success is True
|
|
assert os.path.exists(output_path)
|
|
|
|
# Verify saved image
|
|
saved_img = Image.open(output_path)
|
|
assert saved_img.size == (50, 50)
|
|
|
|
|
|
def test_create_mask_from_alpha():
|
|
"""Test creating mask from alpha channel"""
|
|
with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as f:
|
|
# Create image with alpha channel
|
|
img = Image.new('RGBA', (50, 50))
|
|
# Make half transparent, half opaque
|
|
for x in range(25):
|
|
for y in range(50):
|
|
img.putpixel((x, y), (255, 0, 0, 0)) # Transparent
|
|
for x in range(25, 50):
|
|
for y in range(50):
|
|
img.putpixel((x, y), (0, 255, 0, 255)) # Opaque
|
|
|
|
img.save(f.name, 'PNG')
|
|
|
|
try:
|
|
mask_data = create_mask_from_alpha(f.name)
|
|
|
|
assert mask_data is not None
|
|
|
|
# Verify mask
|
|
mask_img = Image.open(io.BytesIO(mask_data))
|
|
assert mask_img.mode == 'L' # Grayscale
|
|
assert mask_img.size == (50, 50)
|
|
finally:
|
|
os.unlink(f.name)
|
|
|
|
|
|
def test_base64_encoding_decoding():
|
|
"""Test base64 encoding and decoding"""
|
|
# Create test data
|
|
original_data = b"Test image data bytes"
|
|
|
|
# Encode
|
|
encoded = encode_image_base64(original_data)
|
|
assert isinstance(encoded, str)
|
|
|
|
# Decode
|
|
decoded = decode_image_base64(encoded)
|
|
assert decoded == original_data
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-v"])
|