test_startup.py

#!/usr/bin/env python3
"""
Quick startup test for DReamMachine
Tests if all modules can be loaded without crashing
"""

import os
import sys

def test_imports():
    """Test if all modules can be imported"""
    print("Testing imports...")

    try:
        import yaml
        print("✓ yaml imported")
    except Exception as e:
        print(f"✗ yaml failed: {e}")
        return False

    try:
        import gradio
        print("✓ gradio imported")
    except Exception as e:
        print(f"✗ gradio failed: {e}")
        return False

    try:
        from huggingface_hub import InferenceClient
        print("✓ huggingface_hub imported")
    except Exception as e:
        print(f"✗ huggingface_hub failed: {e}")
        return False

    try:
        from datasets import Dataset
        print("✓ datasets imported")
    except Exception as e:
        print(f"✗ datasets failed: {e}")
        return False

    return True


def test_config():
    """Test if config.yaml can be loaded"""
    print("\nTesting config.yaml...")

    try:
        import yaml
        with open('config.yaml', 'r') as f:
            config = yaml.safe_load(f)
        print("✓ config.yaml loaded successfully")
        return True
    except Exception as e:
        print(f"✗ config.yaml failed: {e}")
        return False


def test_modules():
    """Test if custom modules can be imported"""
    print("\nTesting custom modules...")

    try:
        from prompt_manager import PromptManager
        print("✓ prompt_manager imported")
    except Exception as e:
        print(f"✗ prompt_manager failed: {e}")
        return False

    try:
        from llm_agent import LLMAgent
        print("✓ llm_agent imported")
    except Exception as e:
        print(f"✗ llm_agent failed: {e}")
        return False

    try:
        from data_logger import DataLogger
        print("✓ data_logger imported")
    except Exception as e:
        print(f"✗ data_logger failed: {e}")
        return False

    try:
        from orchestrator import DreamOrchestrator
        print("✓ orchestrator imported")
    except Exception as e:
        print(f"✗ orchestrator failed: {e}")
        return False

    return True


def test_hf_token():
    """Test if HF_TOKEN is set"""
    print("\nTesting HF_TOKEN...")

    token = os.getenv('HF_TOKEN')
    if token:
        print(f"✓ HF_TOKEN is set (length: {len(token)})")
        return True
    else:
        print("✗ HF_TOKEN is NOT set")
        print("  Set it with: export HF_TOKEN=your_token_here")
        return False


def main():
    """Run all tests"""
    print("=" * 60)
    print("DReamMachine Startup Test")
    print("=" * 60)

    results = []

    results.append(("Imports", test_imports()))
    results.append(("Config", test_config()))
    results.append(("Modules", test_modules()))
    results.append(("HF Token", test_hf_token()))

    print("\n" + "=" * 60)
    print("Test Results:")
    print("=" * 60)

    for name, passed in results:
        status = "✓ PASS" if passed else "✗ FAIL"
        print(f"{status} - {name}")

    all_passed = all(r[1] for r in results)

    print("\n" + "=" * 60)
    if all_passed:
        print("✓ All tests passed! Ready to run DReamMachine.")
        print("\nStart the app with:")
        print("  python app.py")
    else:
        print("✗ Some tests failed. Please fix the issues above.")
    print("=" * 60)

    return 0 if all_passed else 1


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