198 lines
6.9 KiB
Python
198 lines
6.9 KiB
Python
"""Unit tests for image utilities"""
|
|
|
|
import unittest
|
|
import tempfile
|
|
import base64
|
|
from pathlib import Path
|
|
from PIL import Image
|
|
import io
|
|
|
|
# Add src to path
|
|
import sys
|
|
sys.path.insert(0, str(Path(__file__).parent.parent / 'src'))
|
|
|
|
from src.utils.image_utils import (
|
|
get_file_size_mb,
|
|
validate_image_file,
|
|
get_image_dimensions,
|
|
get_image_dimensions_from_bytes,
|
|
encode_image_base64,
|
|
decode_image_base64,
|
|
save_image,
|
|
get_optimal_aspect_ratio,
|
|
validate_aspect_ratio,
|
|
convert_image_to_base64
|
|
)
|
|
|
|
|
|
class TestImageUtils(unittest.TestCase):
|
|
"""Test cases for image utility functions"""
|
|
|
|
def setUp(self):
|
|
"""Set up test fixtures"""
|
|
self.temp_dir = tempfile.mkdtemp()
|
|
self.temp_path = Path(self.temp_dir)
|
|
|
|
# Create a test image
|
|
self.test_image = Image.new('RGB', (100, 100), color='red')
|
|
self.test_image_path = self.temp_path / 'test_image.png'
|
|
self.test_image.save(self.test_image_path)
|
|
|
|
# Create test image data
|
|
buffer = io.BytesIO()
|
|
self.test_image.save(buffer, format='PNG')
|
|
self.test_image_data = buffer.getvalue()
|
|
|
|
def tearDown(self):
|
|
"""Clean up test fixtures"""
|
|
import shutil
|
|
if self.temp_path.exists():
|
|
shutil.rmtree(self.temp_path)
|
|
|
|
def test_get_file_size_mb(self):
|
|
"""Test file size calculation"""
|
|
size_mb = get_file_size_mb(self.test_image_path)
|
|
self.assertGreater(size_mb, 0)
|
|
self.assertLess(size_mb, 1) # Small test image should be < 1MB
|
|
|
|
def test_validate_image_file_success(self):
|
|
"""Test successful image validation"""
|
|
is_valid, size_mb, error = validate_image_file(str(self.test_image_path), 20)
|
|
|
|
self.assertTrue(is_valid)
|
|
self.assertGreater(size_mb, 0)
|
|
self.assertIsNone(error)
|
|
|
|
def test_validate_image_file_not_found(self):
|
|
"""Test validation of non-existent file"""
|
|
is_valid, size_mb, error = validate_image_file('nonexistent.png', 20)
|
|
|
|
self.assertFalse(is_valid)
|
|
self.assertEqual(size_mb, 0)
|
|
self.assertIn('File not found', error)
|
|
|
|
def test_validate_image_file_too_large(self):
|
|
"""Test validation of file too large"""
|
|
# Test with very small limit
|
|
is_valid, size_mb, error = validate_image_file(str(self.test_image_path), 0.001)
|
|
|
|
self.assertFalse(is_valid)
|
|
self.assertIn('exceeds', error)
|
|
|
|
def test_get_image_dimensions(self):
|
|
"""Test getting image dimensions"""
|
|
width, height = get_image_dimensions(str(self.test_image_path))
|
|
self.assertEqual(width, 100)
|
|
self.assertEqual(height, 100)
|
|
|
|
def test_get_image_dimensions_from_bytes(self):
|
|
"""Test getting dimensions from image bytes"""
|
|
width, height = get_image_dimensions_from_bytes(self.test_image_data)
|
|
self.assertEqual(width, 100)
|
|
self.assertEqual(height, 100)
|
|
|
|
def test_encode_decode_base64(self):
|
|
"""Test base64 encoding and decoding"""
|
|
# Encode
|
|
b64_string = encode_image_base64(self.test_image_data)
|
|
self.assertIsInstance(b64_string, str)
|
|
|
|
# Decode
|
|
decoded_data = decode_image_base64(b64_string)
|
|
self.assertEqual(decoded_data, self.test_image_data)
|
|
|
|
def test_decode_base64_with_data_url(self):
|
|
"""Test decoding base64 with data URL prefix"""
|
|
b64_string = encode_image_base64(self.test_image_data)
|
|
data_url = f"data:image/png;base64,{b64_string}"
|
|
|
|
decoded_data = decode_image_base64(data_url)
|
|
self.assertEqual(decoded_data, self.test_image_data)
|
|
|
|
def test_save_image(self):
|
|
"""Test saving image data to file"""
|
|
output_path = self.temp_path / 'output.png'
|
|
|
|
success = save_image(self.test_image_data, str(output_path))
|
|
self.assertTrue(success)
|
|
self.assertTrue(output_path.exists())
|
|
|
|
# Verify file content
|
|
with open(output_path, 'rb') as f:
|
|
saved_data = f.read()
|
|
self.assertEqual(saved_data, self.test_image_data)
|
|
|
|
def test_get_optimal_aspect_ratio(self):
|
|
"""Test optimal aspect ratio calculation"""
|
|
# Test square image
|
|
ratio = get_optimal_aspect_ratio(100, 100)
|
|
self.assertEqual(ratio, "1:1")
|
|
|
|
# Test wide image
|
|
ratio = get_optimal_aspect_ratio(160, 90)
|
|
self.assertEqual(ratio, "16:9")
|
|
|
|
# Test tall image
|
|
ratio = get_optimal_aspect_ratio(90, 160)
|
|
self.assertEqual(ratio, "9:16")
|
|
|
|
def test_validate_aspect_ratio(self):
|
|
"""Test aspect ratio validation"""
|
|
# Test matching ratio
|
|
self.assertTrue(validate_aspect_ratio(100, 100, "1:1"))
|
|
self.assertTrue(validate_aspect_ratio(160, 90, "16:9"))
|
|
|
|
# Test non-matching ratio (within tolerance)
|
|
self.assertTrue(validate_aspect_ratio(161, 90, "16:9")) # Small difference
|
|
|
|
# Test non-matching ratio (outside tolerance)
|
|
self.assertFalse(validate_aspect_ratio(200, 100, "1:1"))
|
|
|
|
def test_convert_image_to_base64(self):
|
|
"""Test converting image file to base64"""
|
|
b64_string = convert_image_to_base64(str(self.test_image_path))
|
|
|
|
self.assertIsInstance(b64_string, str)
|
|
|
|
# Verify we can decode it back
|
|
decoded_data = decode_image_base64(b64_string)
|
|
|
|
# Images should have same dimensions
|
|
width, height = get_image_dimensions_from_bytes(decoded_data)
|
|
self.assertEqual(width, 100)
|
|
self.assertEqual(height, 100)
|
|
|
|
def create_large_image_file(self, size_mb: float) -> Path:
|
|
"""Helper to create a large image file for testing"""
|
|
# Calculate dimensions for target size (rough estimate)
|
|
# PNG compression varies, so this is approximate
|
|
pixels = int((size_mb * 1024 * 1024) / 4) # 4 bytes per pixel (RGBA)
|
|
dimension = int(pixels ** 0.5)
|
|
|
|
large_image = Image.new('RGBA', (dimension, dimension), color='red')
|
|
large_image_path = self.temp_path / 'large_image.png'
|
|
large_image.save(large_image_path)
|
|
|
|
return large_image_path
|
|
|
|
def test_large_image_handling(self):
|
|
"""Test handling of large images"""
|
|
# This test might be slow, so we'll use a smaller "large" image
|
|
try:
|
|
large_path = self.create_large_image_file(0.1) # 0.1 MB
|
|
|
|
# Test validation
|
|
is_valid, size_mb, error = validate_image_file(str(large_path), 20)
|
|
self.assertTrue(is_valid)
|
|
|
|
# Test conversion to base64
|
|
b64_string = convert_image_to_base64(str(large_path))
|
|
self.assertIsInstance(b64_string, str)
|
|
|
|
except Exception as e:
|
|
self.skipTest(f"Large image test skipped due to: {e}")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|