trax/test_database_setup.py

#!/usr/bin/env python3
"""Test script to verify database setup and models.

This script tests the database connection, models, and basic operations
to ensure everything is working correctly.
"""

import sys
from pathlib import Path

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

from src.config import config
from src.database import Base, get_registered_models
from src.database.connection import test_connection, get_connection_info
from src.database.utils import create_tables, get_table_info


def test_configuration():
    """Test database configuration."""
    print("🔧 Testing Database Configuration")
    print("=" * 40)
    
    # Check if DATABASE_URL is configured
    if not config.DATABASE_URL:
        print("❌ DATABASE_URL not configured")
        return False
    
    print(f"✅ DATABASE_URL: {config.DATABASE_URL}")
    
    # Display connection info
    conn_info = get_connection_info()
    print(f"✅ Host: {conn_info['host']}")
    print(f"✅ Database: {conn_info['database']}")
    print(f"✅ Pool Size: {conn_info['pool_size']}")
    print(f"✅ Max Overflow: {conn_info['max_overflow']}")
    
    return True


def test_models():
    """Test model registration."""
    print("\n📋 Testing Model Registration")
    print("=" * 40)
    
    # Get registered models
    models = get_registered_models()
    
    if not models:
        print("❌ No models registered")
        return False
    
    print(f"✅ Found {len(models)} registered models:")
    for model_name in models:
        print(f"   - {model_name}")
    
    # Check for required models
    required_models = ["MediaFile", "TranscriptionJob", "TranscriptionResult", "ProcessingJob"]
    missing_models = [name for name in required_models if name not in models]
    
    if missing_models:
        print(f"❌ Missing required models: {missing_models}")
        return False
    
    print("✅ All required models are registered")
    return True


def test_database_connection():
    """Test database connection."""
    print("\n🔌 Testing Database Connection")
    print("=" * 40)
    
    try:
        if test_connection():
            print("✅ Database connection successful")
            return True
        else:
            print("❌ Database connection failed")
            return False
    except Exception as e:
        print(f"❌ Database connection error: {e}")
        return False


def test_table_creation():
    """Test table creation."""
    print("\n📊 Testing Table Creation")
    print("=" * 40)
    
    try:
        # Create tables
        create_tables()
        print("✅ Tables created successfully")
        
        # Get table info
        table_info = get_table_info()
        print(f"✅ Found {len(table_info)} tables:")
        
        for table_name, info in table_info.items():
            print(f"   - {table_name}: {len(info['columns'])} columns")
        
        return True
    except Exception as e:
        print(f"❌ Table creation error: {e}")
        return False


def test_jsonb_functionality():
    """Test JSONB functionality."""
    print("\n🔍 Testing JSONB Functionality")
    print("=" * 40)
    
    try:
        from src.database.connection import get_db_session
        from src.database.models import MediaFile
        
        with get_db_session() as session:
            # Test JSONB insert
            test_metadata = {
                "quality": "high",
                "format": "mp4",
                "duration": 120.5,
                "tags": ["test", "sample"]
            }
            
            # Create a test media file
            media_file = MediaFile(
                filename="test_file.mp4",
                file_size=1024000,
                duration=120.5,
                source_path="/path/to/test_file.mp4",
                file_metadata=test_metadata
            )
            
            session.add(media_file)
            session.commit()
            
            print("✅ JSONB insert successful")
            
            # Test JSONB query
            from src.database.utils import jsonb_contains
            
            # Find files with specific metadata
            results = jsonb_contains(session, MediaFile, "file_metadata", {"quality": "high"})
            
            if results:
                print(f"✅ JSONB query successful: found {len(results)} results")
                
                # Clean up test data
                session.delete(media_file)
                session.commit()
                print("✅ Test data cleaned up")
            else:
                print("❌ JSONB query failed")
                return False
            
            return True
            
    except Exception as e:
        print(f"❌ JSONB functionality error: {e}")
        return False


def main():
    """Run all database tests."""
    print("🚀 Database Setup Test Suite")
    print("=" * 50)
    
    tests = [
        ("Configuration", test_configuration),
        ("Models", test_models),
        ("Connection", test_database_connection),
        ("Tables", test_table_creation),
        ("JSONB", test_jsonb_functionality),
    ]
    
    results = []
    
    for test_name, test_func in tests:
        try:
            result = test_func()
            results.append((test_name, result))
        except Exception as e:
            print(f"❌ {test_name} test failed with exception: {e}")
            results.append((test_name, False))
    
    # Summary
    print("\n📊 Test Summary")
    print("=" * 50)
    
    passed = 0
    total = len(results)
    
    for test_name, result in results:
        status = "✅ PASS" if result else "❌ FAIL"
        print(f"{status} {test_name}")
        if result:
            passed += 1
    
    print(f"\nOverall: {passed}/{total} tests passed")
    
    if passed == total:
        print("🎉 All tests passed! Database setup is working correctly.")
        return 0
    else:
        print("⚠️  Some tests failed. Please check the configuration.")
        return 1


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