test_base64_image_transfer.py

#!/usr/bin/env python3
"""
Test Base64 Image Transfer
==========================

Test script to verify base64 image transfer and reconstruction.

Usage:
    python3 test_base64_image_transfer.py [image_path]
"""

import sys
import os
import base64
from bluetooth_image_base64 import BluetoothImageBase64Transfer

def create_test_image():
    """Create a simple test image if PIL is available"""
    try:
        from PIL import Image, ImageDraw
        
        # Create a simple 100x100 test image
        img = Image.new('RGB', (100, 100), color='lightblue')
        draw = ImageDraw.Draw(img)
        
        # Draw some text and shapes
        draw.text((10, 40), "TEST", fill='black')
        draw.rectangle([10, 10, 90, 90], outline='red', width=3)
        
        test_path = "test_base64_image.png"
        img.save(test_path)
        print(f"✅ Created test image: {test_path}")
        return test_path
        
    except ImportError:
        print("❌ PIL not available, cannot create test image")
        return None

def image_to_base64(image_path):
    """Convert image to base64 string"""
    try:
        with open(image_path, 'rb') as f:
            image_data = f.read()
        base64_data = base64.b64encode(image_data).decode('utf-8')
        return base64_data
    except Exception as e:
        print(f"❌ Error converting image to base64: {e}")
        return None

def base64_to_image(base64_data, output_path):
    """Convert base64 string back to image"""
    try:
        # Remove data URL prefix if present
        if base64_data.startswith('data:'):
            base64_data = base64_data.split(',', 1)[1]
        
        image_data = base64.b64decode(base64_data)
        with open(output_path, 'wb') as f:
            f.write(image_data)
        return True
    except Exception as e:
        print(f"❌ Error converting base64 to image: {e}")
        return False

def test_base64_transfer():
    """Test the base64 image transfer protocol"""
    print("🧪 BASE64 IMAGE TRANSFER TEST")
    print("=" * 50)
    
    # Get image path from command line or create test image
    if len(sys.argv) > 1:
        image_path = sys.argv[1]
    else:
        image_path = create_test_image()
        if not image_path:
            print("❌ No test image available. Please provide an image path.")
            return
    
    if not os.path.exists(image_path):
        print(f"❌ Image not found: {image_path}")
        return
    
    print(f"📸 Testing with image: {image_path}")
    print(f"📊 Image size: {os.path.getsize(image_path)} bytes")
    
    # Convert image to base64
    print(f"\n🔧 Converting image to base64...")
    base64_data = image_to_base64(image_path)
    if not base64_data:
        print(f"❌ Failed to convert image to base64")
        return
    
    print(f"✅ Base64 conversion successful")
    print(f"   📊 Base64 size: {len(base64_data)} characters")
    print(f"   📋 First 50 chars: {base64_data[:50]}...")
    print(f"   📋 Last 50 chars: ...{base64_data[-50:]}")
    
    # Test the protocol
    sender_transfer = BluetoothImageBase64Transfer(max_chunk_size=500)
    receiver_transfer = BluetoothImageBase64Transfer(max_chunk_size=500)
    
    sender_name = "TestSender"
    receiver_name = "TestReceiver"
    filename = os.path.basename(image_path)
    
    try:
        print(f"\n📤 Preparing base64 data for transfer...")
        message_chunks = sender_transfer.prepare_base64_for_transfer(
            base64_data, filename, sender_name, receiver_name
        )
        
        print(f"✅ Prepared {len(message_chunks)} message chunks")
        
        # Simulate sending and receiving
        print(f"\n📡 Simulating Bluetooth transfer...")
        received_base64 = None
        
        for i, chunk in enumerate(message_chunks):
            print(f"   📦 Processing chunk {i+1}/{len(message_chunks)}")
            
            # Show chunk details for first few and last few
            if i < 3 or i >= len(message_chunks) - 3:
                print(f"      📋 Chunk content: {chunk[:100]}...")
            
            # Simulate receiving the message
            result = receiver_transfer.process_received_message(chunk)
            
            if result and result.get('type') == 'image_complete':
                print(f"\n✅ Base64 image transfer complete!")
                
                # Get the reconstructed base64 data
                filename, received_base64 = receiver_transfer.get_completed_image_base64(result['msg_id'])
                break
        
        if received_base64:
            print(f"\n🔍 Verifying reconstructed base64...")
            print(f"   📊 Original size: {len(base64_data)} characters")
            print(f"   📊 Received size: {len(received_base64)} characters")
            print(f"   📋 First 50 chars: {received_base64[:50]}...")
            print(f"   📋 Last 50 chars: ...{received_base64[-50:]}")
            
            # Compare the data
            if base64_data == received_base64:
                print(f"   ✅ Base64 data matches perfectly!")
            else:
                print(f"   ❌ Base64 data mismatch!")
                
                # Find where they differ
                for i, (orig, recv) in enumerate(zip(base64_data, received_base64)):
                    if orig != recv:
                        print(f"   🔍 First difference at position {i}: '{orig}' vs '{recv}'")
                        break
                
                if len(base64_data) != len(received_base64):
                    print(f"   🔍 Length difference: original {len(base64_data)}, received {len(received_base64)}")
            
            # Try to reconstruct the image
            output_path = f"received_{filename}"
            print(f"\n💾 Saving reconstructed image as: {output_path}")
            
            if base64_to_image(received_base64, output_path):
                print(f"✅ Successfully saved reconstructed image")
                
                # Verify file size
                original_size = os.path.getsize(image_path)
                received_size = os.path.getsize(output_path)
                
                print(f"📊 Original file size: {original_size} bytes")
                print(f"📊 Received file size: {received_size} bytes")
                
                if original_size == received_size:
                    print(f"✅ File sizes match!")
                else:
                    print(f"❌ File size mismatch!")
            else:
                print(f"❌ Failed to save reconstructed image")
        else:
            print(f"❌ No base64 data received")
        
        print(f"\n🎉 Test completed!")
        
    except Exception as e:
        print(f"❌ Test failed: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    test_base64_transfer()