test_database_migration.py

#!/usr/bin/env python3
"""
Test script for the updated database.py functions to ensure they work with Supabase.
Run this script to verify the migration from Prisma to Supabase was successful.
"""

import asyncio
import sys
import os

# Add the api directory to the path
sys.path.append(os.path.join(os.path.dirname(__file__), "api"))

from database import (
    get_supabase_client,
    get_pg_connection,
    get_all_videos,
    get_all_segments,
    search_segments_by_text,
    get_search_statistics,
    close_connections,
)


async def test_database_functions():
    """Test all database functions to ensure they work with Supabase."""
    print("🧪 Testing Updated Database Functions")
    print("=" * 50)

    try:
        # Test Supabase client connection
        print("1. Testing Supabase client connection...")
        client = get_supabase_client()
        print("   ✅ Supabase client initialized successfully")

        # Test PostgreSQL connection
        print("\n2. Testing PostgreSQL connection...")
        try:
            connection = await get_pg_connection()
            print("   ✅ PostgreSQL connection established successfully")
        except Exception as e:
            print(f"   ⚠️  PostgreSQL connection failed (fallback will work): {e}")

        # Test get_all_videos
        print("\n3. Testing get_all_videos()...")
        videos = await get_all_videos()
        print(f"   ✅ Retrieved {len(videos)} videos")
        if videos:
            video = videos[0]
            print(
                f"   📹 Sample video: {video.get('title', 'Unknown')} (ID: {video.get('id', 'Unknown')})"
            )

        # Test get_all_segments
        print("\n4. Testing get_all_segments()...")
        segments = await get_all_segments()
        print(f"   ✅ Retrieved {len(segments)} segments")
        if segments:
            segment = segments[0]
            print(
                f"   📝 Sample segment: '{segment.get('text', '')[:50]}...' (ID: {segment.get('id', 'Unknown')})"
            )

        # Test text search
        print("\n5. Testing search_segments_by_text()...")
        if segments:
            # Use a word from the first segment as a test query
            first_segment_text = segments[0].get("text", "")
            if first_segment_text:
                test_word = (
                    first_segment_text.split()[0]
                    if first_segment_text.split()
                    else "test"
                )
                search_results = await search_segments_by_text(test_word, limit=5)
                print(
                    f"   ✅ Text search for '{test_word}' returned {len(search_results)} results"
                )
            else:
                print("   ⚠️  No text content found to test search")
        else:
            print("   ⚠️  No segments available to test search")

        # Test statistics
        print("\n6. Testing get_search_statistics()...")
        stats = await get_search_statistics()
        print(f"   ✅ Statistics retrieved:")
        for key, value in stats.items():
            print(f"      {key}: {value}")

        print("\n🎉 All tests completed successfully!")
        print(
            "✅ Database migration from Prisma to Supabase appears to be working correctly."
        )

    except Exception as e:
        print(f"\n❌ Test failed with error: {e}")
        import traceback

        traceback.print_exc()
        return False

    finally:
        # Clean up connections
        await close_connections()
        print("\n🧹 Cleaned up database connections")

    return True


async def main():
    """Main function to run the tests."""
    print("🚀 Starting Database Migration Test")
    print("-" * 50)

    success = await test_database_functions()

    if success:
        print("\n✅ Migration test PASSED! The database.py file is ready to use.")
        sys.exit(0)
    else:
        print("\n❌ Migration test FAILED! Please check the errors above.")
        sys.exit(1)


if __name__ == "__main__":
    asyncio.run(main())