scratch/test_db.py

import asyncio
import sys
from pathlib import Path
from datetime import datetime, timezone

# Add project root to path
sys.path.append(str(Path(__file__).resolve().parent.parent))

async def test_db_operations():
    from database.db import init_db, AsyncSessionLocal, engine
    from database.models import Alpha
    from sqlalchemy import select

    # Initialize tables
    print("Initializing database...")
    await init_db()

    # Open session
    print("Opening async session...")
    async with AsyncSessionLocal() as session:
        # Create a new Alpha record
        new_alpha = Alpha(
            expression="rank(close - open)",
            language="FASTEXPR",
            region="USA",
            universe="TOP3000",
            delay=1,
            decay=4,
            truncation=0.08,
            neutralization="SUBINDUSTRY",
            pasteurization="On",
            status="PENDING_GEN",
            retry_count=0
        )
        session.add(new_alpha)
        await session.commit()
        print(f"Committed new Alpha with expression: {new_alpha.expression}")

        # Query the Alpha
        result = await session.execute(select(Alpha).where(Alpha.expression == "rank(close - open)"))
        alpha_record = result.scalars().first()
        assert alpha_record is not None
        print(f"Queried Alpha back: id={alpha_record.id}, region={alpha_record.region}, status={alpha_record.status}")
        
        # Update Alpha status
        alpha_record.status = "SIMULATING"
        await session.commit()
        print(f"Updated Alpha status to SIMULATING. New status: {alpha_record.status}")
        
        # Clean up by deleting the alpha
        await session.delete(alpha_record)
        await session.commit()
        print("Cleaned up and deleted test Alpha record.")
        
    await engine.dispose()
    print("Database connection pool disposed.")

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