Lyra/test_database_connections.py

"""
Test database connections for Lyra AI.

This script tests both PostgreSQL and Redis connections to ensure
they're properly configured before running Lyra.
"""

import asyncio
import sys
import os
from pathlib import Path

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

async def test_postgresql():
    """Test PostgreSQL connection."""
    print("Testing PostgreSQL connection...")

    try:
        import asyncpg
        from lyra.config import config

        # Parse the database URL
        database_url = config.database_url
        print(f"Connecting to: {database_url.replace('your_password_here', '****')}")

        # Test connection
        conn = await asyncpg.connect(database_url)

        # Test query
        result = await conn.fetchval('SELECT version()')
        print(f"✅ PostgreSQL connected successfully!")
        print(f"   Version: {result}")

        # Test creating a simple table
        await conn.execute('''
            CREATE TABLE IF NOT EXISTS test_connection (
                id SERIAL PRIMARY KEY,
                message TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        ''')

        # Insert test data
        await conn.execute(
            'INSERT INTO test_connection (message) VALUES ($1)',
            'Lyra database test successful'
        )

        # Verify test data
        test_result = await conn.fetchval(
            'SELECT message FROM test_connection ORDER BY id DESC LIMIT 1'
        )
        print(f"   Test query result: {test_result}")

        # Clean up test table
        await conn.execute('DROP TABLE IF EXISTS test_connection')

        await conn.close()
        return True

    except ImportError:
        print("❌ asyncpg not installed. Run: pip install asyncpg")
        return False
    except Exception as e:
        print(f"❌ PostgreSQL connection failed: {e}")
        print("\n💡 Troubleshooting tips:")
        print("   1. Make sure PostgreSQL is running")
        print("   2. Check your password in .env file")
        print("   3. Create the 'lyra' database if it doesn't exist")
        print("   4. Verify PostgreSQL is listening on port 5432")
        return False


async def test_redis():
    """Test Redis connection."""
    print("\nTesting Redis connection...")

    try:
        import redis.asyncio as redis
        from lyra.config import config

        # Connect to Redis
        redis_client = redis.from_url(config.redis_url)

        # Test ping
        response = await redis_client.ping()
        print(f"✅ Redis connected successfully!")
        print(f"   Ping response: {response}")

        # Test basic operations
        await redis_client.set('lyra_test', 'Hello from Lyra AI!')
        test_value = await redis_client.get('lyra_test')
        print(f"   Test value: {test_value.decode('utf-8')}")

        # Clean up
        await redis_client.delete('lyra_test')
        await redis_client.close()

        return True

    except ImportError:
        print("❌ redis not installed. Run: pip install redis")
        return False
    except Exception as e:
        print(f"❌ Redis connection failed: {e}")
        print("\n💡 Troubleshooting tips:")
        print("   1. Make sure Redis/Memurai is running")
        print("   2. Check if port 6379 is available")
        print("   3. Try restarting the Redis service")
        return False


async def create_lyra_database():
    """Create the Lyra database if it doesn't exist."""
    print("\nCreating Lyra database...")

    try:
        import asyncpg
        from lyra.config import config

        # Connect to postgres database (default)
        base_url = config.database_url.replace('/lyra', '/postgres')
        conn = await asyncpg.connect(base_url)

        # Check if lyra database exists
        db_exists = await conn.fetchval(
            "SELECT 1 FROM pg_database WHERE datname = 'lyra'"
        )

        if not db_exists:
            await conn.execute('CREATE DATABASE lyra')
            print("✅ Created 'lyra' database")
        else:
            print("✅ 'lyra' database already exists")

        await conn.close()
        return True

    except Exception as e:
        print(f"❌ Failed to create database: {e}")
        return False


async def main():
    """Main test function."""
    print("=" * 60)
    print("LYRA AI - DATABASE CONNECTION TESTS")
    print("=" * 60)

    # Load environment variables
    try:
        from dotenv import load_dotenv
        load_dotenv()
        print("✅ Environment variables loaded")
    except ImportError:
        print("⚠️  python-dotenv not found - using system environment")

    # Check if database URL is configured
    if 'your_password_here' in os.getenv('DATABASE_URL', ''):
        print("\n❌ Please update your PostgreSQL password in .env file")
        print("   Replace 'your_password_here' with your actual PostgreSQL password")
        return False

    # Test database creation
    await create_lyra_database()

    # Test connections
    postgresql_ok = await test_postgresql()
    redis_ok = await test_redis()

    print("\n" + "=" * 60)
    if postgresql_ok and redis_ok:
        print("🎉 ALL DATABASE TESTS PASSED!")
        print("Lyra is ready for full deployment with database support.")
        print("\nNext step: Run 'python -m lyra.main' to start Lyra")
    else:
        print("❌ SOME TESTS FAILED")
        print("Please fix the issues above before running Lyra")
    print("=" * 60)

    return postgresql_ok and redis_ok


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