gpt-edit/tests/test_webp_optimization.py

"""
Test script for WebP auto-optimization feature
"""

import sys
import os
from pathlib import Path
import tempfile

# Add project root to path
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

def test_webp_optimization():
    """Test WebP optimization functionality"""
    print("\n" + "="*60)
    print("Testing WebP Auto-Optimization Feature")
    print("="*60)
    
    try:
        from src.utils.image_utils import (
            optimize_image_to_size_limit,
            convert_to_png_with_size_limit,
            get_file_size_mb
        )
        print("✓ Image utils imported successfully")
        
        # Create a test image
        from PIL import Image
        import io
        
        # Create a large test image (simulate >4MB)
        print("\nCreating test image...")
        img = Image.new('RGBA', (2048, 2048), color=(255, 0, 0, 255))
        
        # Save to temp file
        with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp:
            img.save(tmp.name, 'PNG')
            temp_path = tmp.name
        
        print(f"✓ Test image created: {temp_path}")
        original_size = get_file_size_mb(temp_path)
        print(f"  Original size: {original_size:.2f}MB")
        
        # Test WebP conversion
        print("\nTesting WebP conversion...")
        optimized_data, format_used = convert_to_png_with_size_limit(
            temp_path,
            max_size_mb=1.0,  # Force optimization by setting low limit
            prefer_webp=True
        )
        
        optimized_size = len(optimized_data) / (1024 * 1024)
        print(f"✓ Image optimized successfully")
        print(f"  Format used: {format_used}")
        print(f"  Optimized size: {optimized_size:.2f}MB")
        print(f"  Size reduction: {((original_size - optimized_size) / original_size * 100):.1f}%")
        
        # Clean up
        os.unlink(temp_path)
        print("\n✓ Test completed successfully!")
        
        return True
        
    except ImportError as e:
        print(f"✗ Import error: {e}")
        print("  Make sure Pillow is installed: pip install pillow")
        return False
    except Exception as e:
        print(f"✗ Test failed: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_optimization_with_client():
    """Test optimization with OpenAI client"""
    print("\n" + "="*60)
    print("Testing OpenAI Client with Auto-Optimization")
    print("="*60)
    
    try:
        from src.connector.config import Config
        from src.connector.openai_client import ImageEditRequest
        
        print("✓ Client modules imported")
        
        # Create test config
        config = Config()
        print("✓ Config created")
        
        # Create test request
        request = ImageEditRequest(
            image_path="test_large_image.png",  # This would be a real file in production
            prompt="Make the image more colorful",
            background="transparent",
            auto_optimize=True  # Enable auto-optimization
        )
        
        print("✓ Request created with auto_optimize=True")
        print("\nIn production, this would:")
        print("  1. Check if image > 4MB")
        print("  2. Automatically convert to WebP if needed")
        print("  3. Find optimal quality setting")
        print("  4. Send optimized image to OpenAI API")
        print("  5. Log optimization details in response")
        
        return True
        
    except Exception as e:
        print(f"✗ Test failed: {e}")
        return False


def test_quality_levels():
    """Test different quality levels for optimization"""
    print("\n" + "="*60)
    print("Testing Quality Levels")
    print("="*60)
    
    try:
        from src.utils.image_utils import optimize_image_to_size_limit
        from PIL import Image
        import tempfile
        
        # Create test image with details
        print("Creating detailed test image...")
        img = Image.new('RGB', (1024, 1024))
        
        # Add some patterns to make compression more interesting
        from PIL import ImageDraw
        draw = ImageDraw.Draw(img)
        for i in range(0, 1024, 20):
            draw.line([(0, i), (1024, i)], fill=(i % 255, 100, 200))
            draw.line([(i, 0), (i, 1024)], fill=(200, i % 255, 100))
        
        # Save to temp file
        with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp:
            img.save(tmp.name, 'PNG')
            temp_path = tmp.name
        
        print(f"✓ Detailed test image created")
        
        # Test different size limits
        size_limits = [0.5, 1.0, 2.0]
        
        for limit in size_limits:
            print(f"\nTesting with {limit}MB limit...")
            try:
                optimized = optimize_image_to_size_limit(
                    temp_path,
                    max_size_mb=limit,
                    format='WEBP'
                )
                size_mb = len(optimized) / (1024 * 1024)
                print(f"  ✓ Optimized to {size_mb:.2f}MB (limit: {limit}MB)")
            except Exception as e:
                print(f"  ✗ Failed: {e}")
        
        # Clean up
        os.unlink(temp_path)
        
        print("\n✓ Quality level tests completed!")
        return True
        
    except Exception as e:
        print(f"✗ Test failed: {e}")
        return False


def main():
    """Run all optimization tests"""
    print("="*60)
    print("GPTEdit WebP Auto-Optimization Test Suite")
    print("="*60)
    
    tests = [
        test_webp_optimization,
        test_optimization_with_client,
        test_quality_levels
    ]
    
    passed = 0
    failed = 0
    
    for test in tests:
        if test():
            passed += 1
        else:
            failed += 1
    
    print("\n" + "="*60)
    print(f"Test Results: {passed} passed, {failed} failed")
    print("="*60)
    
    if failed == 0:
        print("✅ All optimization tests passed!")
        return 0
    else:
        print(f"❌ {failed} test(s) failed")
        return 1


if __name__ == "__main__":
    sys.exit(main())