"""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()