test_pipeline_fix.py

#!/usr/bin/env python3
"""
Test script to verify the pipeline fixes work correctly
"""

import sys
import os
import traceback
from typing import Dict, Any

# Add the project root to the path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

def test_pipeline_fixes():
    """Test the pipeline with improved error handling"""
    
    print("🧪 Testing Monster Generation Pipeline Fixes")
    print("=" * 50)
    
    try:
        # Import the pipeline
        from core.ai_pipeline import MonsterGenerationPipeline
        print("✅ Successfully imported MonsterGenerationPipeline")
        
        # Initialize pipeline
        print("🔧 Initializing pipeline...")
        pipeline = MonsterGenerationPipeline(device="cpu")  # Use CPU for testing
        print("✅ Pipeline initialized successfully")
        
        # Test with a simple text input
        print("\n🚀 Testing monster generation...")
        test_input = "Create a friendly fire monster with wings"
        
        result = pipeline.generate_monster(
            text_input=test_input,
            user_id="test_user"
        )
        
        print(f"\n📊 Generation Result:")
        print(f"Status: {result.get('status', 'unknown')}")
        print(f"Success: {result.get('generation_log', {}).get('success', False)}")
        print(f"Stages completed: {result.get('generation_log', {}).get('stages_completed', [])}")
        print(f"Fallbacks used: {result.get('generation_log', {}).get('fallbacks_used', [])}")
        print(f"Errors: {result.get('generation_log', {}).get('errors', [])}")
        
        if result.get('traits'):
            print(f"Monster name: {result.get('traits', {}).get('name', 'Unknown')}")
            print(f"Monster element: {result.get('traits', {}).get('element', 'Unknown')}")
        
        if result.get('dialogue'):
            print(f"Monster dialogue: {result.get('dialogue', '')}")
        
        print(f"Download files: {result.get('download_files', [])}")
        
        # Clean up
        pipeline.cleanup()
        print("\n🧹 Pipeline cleaned up successfully")
        
        return True
        
    except Exception as e:
        print(f"❌ Test failed with error: {e}")
        print(f"Error type: {type(e).__name__}")
        print("Full traceback:")
        traceback.print_exc()
        return False

def test_fallback_manager():
    """Test the fallback manager"""
    
    print("\n🧪 Testing Fallback Manager")
    print("=" * 30)
    
    try:
        from utils.fallbacks import FallbackManager
        
        fallback = FallbackManager()
        
        # Test text generation fallback
        print("📝 Testing text generation fallback...")
        traits, dialogue = fallback.handle_text_gen_failure("Create a water monster")
        print(f"✅ Generated traits: {traits.get('name', 'Unknown')}")
        print(f"✅ Generated dialogue: {dialogue}")
        
        # Test image generation fallback
        print("🎨 Testing image generation fallback...")
        image = fallback.handle_image_gen_failure("Create a fire monster")
        print(f"✅ Generated image: {type(image)}")
        
        # Test 3D generation fallback
        print("🔲 Testing 3D generation fallback...")
        model_3d = fallback.handle_3d_gen_failure(image)
        print(f"✅ Generated 3D model: {type(model_3d)}")
        
        print("✅ All fallback tests passed!")
        return True
        
    except Exception as e:
        print(f"❌ Fallback test failed: {e}")
        traceback.print_exc()
        return False

def main():
    """Main test function"""
    
    print("🔍 Starting Pipeline Fix Verification")
    print("=" * 50)
    
    # Test fallback manager first (doesn't require heavy models)
    fallback_success = test_fallback_manager()
    
    # Test full pipeline (may fail due to missing models, but should show better error handling)
    pipeline_success = test_pipeline_fixes()
    
    print("\n" + "=" * 50)
    print("📋 Test Results Summary:")
    print(f"Fallback Manager: {'✅ PASSED' if fallback_success else '❌ FAILED'}")
    print(f"Pipeline: {'✅ PASSED' if pipeline_success else '❌ FAILED'}")
    
    if fallback_success and pipeline_success:
        print("\n🎉 All tests passed! Pipeline fixes are working correctly.")
    elif fallback_success:
        print("\n⚠️ Fallback manager works, but pipeline may need model dependencies.")
        print("This is expected if models aren't installed.")
    else:
        print("\n❌ Some tests failed. Check the error messages above.")
    
    return fallback_success and pipeline_success

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)