test_supabase_upload.py

#!/usr/bin/env python3
"""
Test script for Supabase Storage upload functionality.

Usage:
    python test_supabase_upload.py

Make sure to set environment variables before running:
    export SUPABASE_URL="https://your-project-id.supabase.co"
    export SUPABASE_KEY="your-service-role-key"
    export SUPABASE_BUCKET="outputs"  # optional
"""

import os
import sys
import tempfile
import time


def create_test_file(size_mb: int = 10) -> str:
    """Create a temporary test file of specified size in MB."""
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.bin')
    
    # Write random data
    chunk_size = 1024 * 1024  # 1 MB
    for _ in range(size_mb):
        temp_file.write(os.urandom(chunk_size))
    
    temp_file.close()
    print(f"Created test file: {temp_file.name} ({size_mb} MB)")
    return temp_file.name


def main():
    # Import the upload function from app.py
    try:
        # Add parent directory to path
        sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
        from app import upload_to_supabase_storage
    except ImportError as e:
        print(f"Error importing upload function: {e}")
        print("Make sure app.py is in the same directory.")
        return 1

    # Check environment variables
    supabase_url = os.environ.get("SUPABASE_URL")
    supabase_key = os.environ.get("SUPABASE_KEY")
    supabase_bucket = os.environ.get("SUPABASE_BUCKET", "outputs")

    if not supabase_url:
        print("ERROR: SUPABASE_URL environment variable not set")
        print("Set it with: export SUPABASE_URL='https://your-project-id.supabase.co'")
        return 1

    if not supabase_key:
        print("ERROR: SUPABASE_KEY environment variable not set")
        print("Set it with: export SUPABASE_KEY='your-service-role-key'")
        return 1

    print("=" * 80)
    print("Supabase Storage Upload Test")
    print("=" * 80)
    print(f"URL: {supabase_url}")
    print(f"Bucket: {supabase_bucket}")
    print(f"Key: {supabase_key[:20]}... (truncated)")
    print("=" * 80)

    # Test 1: Small file upload (1 MB)
    print("\nTest 1: Small file upload (1 MB)")
    print("-" * 80)
    test_file_small = create_test_file(1)
    
    try:
        start_time = time.time()
        url = upload_to_supabase_storage(
            file_path=test_file_small,
            remote_path="test/small_file.bin",
            supabase_url=supabase_url,
            supabase_key=supabase_key,
            bucket_name=supabase_bucket
        )
        elapsed = time.time() - start_time
        
        if url:
            print(f"✅ SUCCESS: Uploaded in {elapsed:.2f} seconds")
            print(f"URL: {url}")
        else:
            print("❌ FAILED: Upload returned None")
    except Exception as e:
        print(f"❌ ERROR: {e}")
    finally:
        os.unlink(test_file_small)

    # Test 2: Medium file upload (50 MB)
    print("\nTest 2: Medium file upload (50 MB)")
    print("-" * 80)
    test_file_medium = create_test_file(50)
    
    try:
        start_time = time.time()
        url = upload_to_supabase_storage(
            file_path=test_file_medium,
            remote_path="test/medium_file.bin",
            supabase_url=supabase_url,
            supabase_key=supabase_key,
            bucket_name=supabase_bucket
        )
        elapsed = time.time() - start_time
        
        if url:
            print(f"✅ SUCCESS: Uploaded in {elapsed:.2f} seconds")
            print(f"URL: {url}")
        else:
            print("❌ FAILED: Upload returned None")
    except Exception as e:
        print(f"❌ ERROR: {e}")
    finally:
        os.unlink(test_file_medium)

    # Test 3: Large file upload (200 MB) - optional
    print("\nTest 3: Large file upload (200 MB) - Optional")
    print("-" * 80)
    response = input("Run large file test? This will create and upload a 200MB file (y/N): ")
    
    if response.lower() == 'y':
        test_file_large = create_test_file(200)
        
        try:
            start_time = time.time()
            url = upload_to_supabase_storage(
                file_path=test_file_large,
                remote_path="test/large_file.bin",
                supabase_url=supabase_url,
                supabase_key=supabase_key,
                bucket_name=supabase_bucket
            )
            elapsed = time.time() - start_time
            
            if url:
                print(f"✅ SUCCESS: Uploaded in {elapsed:.2f} seconds")
                print(f"Upload speed: {200 / elapsed:.2f} MB/s")
                print(f"URL: {url}")
            else:
                print("❌ FAILED: Upload returned None")
        except Exception as e:
            print(f"❌ ERROR: {e}")
        finally:
            os.unlink(test_file_large)
    else:
        print("Skipped large file test")

    print("\n" + "=" * 80)
    print("Testing complete!")
    print("=" * 80)
    
    return 0


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