test_mesh_generation.py

"""
Test mesh generation locally before deployment.
Verifies the complete pipeline works end-to-end.
"""

import sys
from pathlib import Path

# Add parent directory to path
sys.path.insert(0, str(Path(__file__).parent))

from core import AssetPipeline, QUALITY_PRESETS
from utils import MemoryManager


def test_mesh_generation():
    """Test complete mesh generation pipeline."""
    
    print("=" * 80)
    print("🧪 MESH GENERATION TEST")
    print("=" * 80)
    print()
    
    # Initialize
    print("📦 Initializing pipeline...")
    memory_manager = MemoryManager()
    memory_manager.setup_cuda_optimizations()
    
    pipeline = AssetPipeline()
    print("   ✅ Pipeline initialized")
    print()
    
    # Test configuration
    test_cases = [
        {
            "name": "Fast Generation",
            "prompt": "medieval knight with detailed armor",
            "quality": "Fast",
            "expected_time": 45
        },
        {
            "name": "Balanced Generation",
            "prompt": "futuristic mech robot, game asset",
            "quality": "Balanced",
            "expected_time": 60
        }
    ]
    
    results = []
    
    for i, test in enumerate(test_cases, 1):
        print(f"🎯 Test {i}/{len(test_cases)}: {test['name']}")
        print(f"   Prompt: {test['prompt']}")
        print(f"   Quality: {test['quality']}")
        print(f"   Expected time: ~{test['expected_time']}s")
        print()
        
        try:
            # Generate asset
            print("   🚀 Generating...")
            result = pipeline.generate(
                prompt=test['prompt'],
                quality=test['quality']
            )
            
            # Verify result
            if result.glb_path and result.glb_path.exists():
                file_size_mb = result.glb_path.stat().st_size / 1e6
                
                print(f"   ✅ SUCCESS!")
                print(f"   📊 Generation time: {result.metadata.generation_time_s:.1f}s")
                print(f"   📦 File size: {file_size_mb:.2f} MB")
                print(f"   🎨 FLUX steps: {result.metadata.flux_steps}")
                print(f"   🎨 Hunyuan steps: {result.metadata.hunyuan_steps}")
                print(f"   🔧 Optimized: {result.metadata.optimized}")
                print(f"   💾 Cached: {result.cached}")
                print(f"   📁 Output: {result.glb_path}")
                
                results.append({
                    "test": test['name'],
                    "success": True,
                    "time": result.metadata.generation_time_s,
                    "size_mb": file_size_mb,
                    "cached": result.cached
                })
            else:
                print(f"   ❌ FAILED: No output file")
                results.append({
                    "test": test['name'],
                    "success": False,
                    "error": "No output file"
                })
        
        except Exception as e:
            print(f"   ❌ FAILED: {e}")
            results.append({
                "test": test['name'],
                "success": False,
                "error": str(e)
            })
        
        print()
    
    # Summary
    print("=" * 80)
    print("📊 TEST SUMMARY")
    print("=" * 80)
    print()
    
    success_count = sum(1 for r in results if r.get('success', False))
    total_count = len(results)
    
    print(f"✅ Passed: {success_count}/{total_count}")
    print(f"❌ Failed: {total_count - success_count}/{total_count}")
    print()
    
    if success_count == total_count:
        print("🎉 ALL TESTS PASSED!")
        print()
        print("✨ Pipeline is ready for deployment!")
        print()
        print("Next steps:")
        print("1. Run: .\\FINAL_DEPLOYMENT.ps1")
        print("2. Wait 5-10 minutes for Space to build")
        print("3. Test on HF Space with same prompts")
        print()
        return True
    else:
        print("⚠️ SOME TESTS FAILED")
        print()
        print("Failed tests:")
        for r in results:
            if not r.get('success', False):
                print(f"   • {r['test']}: {r.get('error', 'Unknown error')}")
        print()
        print("Fix errors before deployment!")
        print()
        return False


if __name__ == "__main__":
    try:
        success = test_mesh_generation()
        sys.exit(0 if success else 1)
    except KeyboardInterrupt:
        print("\n\n⚠️ Test interrupted by user")
        sys.exit(1)
    except Exception as e:
        print(f"\n\n❌ Test failed with error: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)