scratch/test_generator.py

import asyncio
import sys
import os
from pathlib import Path

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

async def run_tests():
    # 1. Test Ingestion
    print("Testing document ingestion...")
    from utils.ingest_docs import compile_knowledge_context
    context = compile_knowledge_context()
    print("Compiled Context preview:")
    print("-" * 50)
    print(context[:300])
    print("-" * 50)
    assert "DOCUMENT:" in context
    print("Document Ingestion PASSED!\n")

    # 2. Test Prompts
    print("Testing prompt generation...")
    from core.prompts import get_system_prompt, get_user_prompt
    sys_prompt = get_system_prompt()
    user_prompt = get_user_prompt("USA", "TOP3000", "Mean Reversion", "Statistical", context[:100])
    assert "WorldQuant" in sys_prompt
    assert "USA" in user_prompt
    print("System Prompt Preview:")
    print("-" * 50)
    print(sys_prompt[:300])
    print("-" * 50)
    print("Prompt Generation PASSED!\n")

    # 3. Test In-Memory Database Initialization and Generator Parsing
    print("Testing Generator validation and DB queue...")
    from core.generator import AlphaGenerator
    from database.db import init_db
    
    # Initialize DB (WAL, tables etc.)
    await init_db()
    
    # Clear Alpha table for clean test run
    from database.db import AsyncSessionLocal
    from database.models import Alpha
    from sqlalchemy import delete
    async with AsyncSessionLocal() as session:
        await session.execute(delete(Alpha))
        await session.commit()
    print("Database cleared for clean test run.")

    generator = AlphaGenerator()
    
    # Test 3a. Valid JSON Ingestion
    valid_json = """{
        "hypothesis": "Test momentum",
        "expression": "group_neutralize(rank(close - open), subindustry)",
        "region": "USA",
        "universe": "TOP3000",
        "delay": 1,
        "decay": 5,
        "truncation": 0.08,
        "neutralization": "SUBINDUSTRY"
    }"""
    print("Submitting valid candidate JSON to parse_and_queue...")
    result_valid = await generator.parse_and_queue(valid_json)
    assert result_valid is not None
    assert result_valid["region"] == "USA"
    assert result_valid["universe"] == "TOP3000"
    print("Valid Candidate Parsing PASSED!")

    # Test 3b. Invalid region validation
    invalid_region_json = """{
        "hypothesis": "Test bad region",
        "expression": "group_neutralize(rank(close - open), subindustry)",
        "region": "MARS",
        "universe": "TOP3000",
        "delay": 1,
        "decay": 5,
        "truncation": 0.08,
        "neutralization": "SUBINDUSTRY"
    }"""
    print("Submitting candidate with invalid region...")
    result_invalid_region = await generator.parse_and_queue(invalid_region_json)
    assert result_invalid_region is None
    print("Invalid Region Rejection PASSED!")

    # Test 3c. Invalid delay for region (Delay 0 is not allowed for ASI in config.py)
    invalid_delay_json = """{
        "hypothesis": "Test bad delay for ASI",
        "expression": "group_neutralize(rank(close - open), subindustry)",
        "region": "ASI",
        "universe": "ASI3000",
        "delay": 0,
        "decay": 5,
        "truncation": 0.08,
        "neutralization": "SUBINDUSTRY"
    }"""
    print("Submitting candidate with invalid delay for region...")
    result_invalid_delay = await generator.parse_and_queue(invalid_delay_json)
    assert result_invalid_delay is None
    print("Invalid Delay Rejection PASSED!")

    # Test 3d. Fallback Generator JSON Output check
    print("Checking local fallback JSON output...")
    fallback_json = generator._generate_fallback_json()
    result_fallback = await generator.parse_and_queue(fallback_json)
    assert result_fallback is not None
    print(f"Local fallback alpha queued: {result_fallback['expression'][:50]} (Region: {result_fallback['region']})")
    print("Fallback check PASSED!")

    print("\nAll Phase 2 tests completed successfully!")

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