MushroomFleet-ComfyUI-Variationator/tests/test_modular_system.py

"""
Comprehensive test for the modular architecture
Tests all 220 modifiers across 11 variation dimensions
"""

from core.orchestrator import PromptOrchestrator


def test_all_modifier_types():
    """Test that all modifier types load and work correctly"""
    orchestrator = PromptOrchestrator()
    base_prompt = "Modern sports car showcased prominently"
    
    print("="*80)
    print("TESTING MODULAR ARCHITECTURE - ALL 220 MODIFIERS")
    print("="*80)
    
    # Test each modifier type
    modifier_types = [
        ("Lighting", orchestrator.modifiers, orchestrator.apply_lighting_modifier),
        ("Seasonal", orchestrator.seasonal_modifiers, orchestrator.apply_seasonal_modifier),
        ("Location", orchestrator.location_modifiers, orchestrator.apply_location_modifier),
        ("Weather", orchestrator.weather_modifiers, orchestrator.apply_weather_modifier),
        ("Mood", orchestrator.mood_modifiers, orchestrator.apply_mood_modifier),
        ("Camera", orchestrator.camera_modifiers, orchestrator.apply_camera_modifier),
        ("Color Grading", orchestrator.color_grading_modifiers, orchestrator.apply_color_grading_modifier),
        ("Action", orchestrator.action_modifiers, orchestrator.apply_action_modifier),
        ("Era", orchestrator.era_modifiers, orchestrator.apply_era_modifier),
        ("Detail", orchestrator.detail_modifiers, orchestrator.apply_detail_modifier),
        ("Composition", orchestrator.composition_modifiers, orchestrator.apply_composition_modifier),
    ]
    
    total_modifiers = 0
    total_tested = 0
    
    for type_name, modifiers_dict, apply_func in modifier_types:
        print(f"\n{type_name} Modifiers:")
        print("-"*80)
        
        # Test all modifiers in this category
        for mod_id in sorted(modifiers_dict.keys()):
            modifier = modifiers_dict[mod_id]
            
            # Apply the modifier
            try:
                result = apply_func(base_prompt, mod_id, preserve_base=True)
                status = "✓"
                total_tested += 1
            except Exception as e:
                status = f"✗ Error: {e}"
                result = ""
            
            print(f"  {status} #{mod_id}: {modifier.description[:60]}...")
            total_modifiers += 1
        
        print(f"  Loaded: {len(modifiers_dict)} modifiers")
    
    print("\n" + "="*80)
    print(f"SUMMARY: {total_tested}/{total_modifiers} modifiers tested successfully")
    print("="*80)
    
    return total_tested == total_modifiers


def test_wet_dry_distribution():
    """Test wet/dry categorization across applicable modifiers"""
    orchestrator = PromptOrchestrator()
    
    print("\n" + "="*80)
    print("WET/DRY DISTRIBUTION ANALYSIS")
    print("="*80)
    
    categories_with_wetdry = [
        ("Lighting", orchestrator.modifiers),
        ("Seasonal", orchestrator.seasonal_modifiers),
        ("Weather", orchestrator.weather_modifiers),
    ]
    
    for type_name, modifiers_dict in categories_with_wetdry:
        dry_count = sum(1 for m in modifiers_dict.values() if m.category.value == "dry")
        wet_count = sum(1 for m in modifiers_dict.values() if m.category.value == "wet")
        
        print(f"\n{type_name}:")
        print(f"  DRY: {dry_count} modifiers")
        print(f"  WET: {wet_count} modifiers")
        print(f"  Total: {dry_count + wet_count}")


def test_lazy_loading():
    """Test that lazy loading works correctly"""
    print("\n" + "="*80)
    print("LAZY LOADING TEST")
    print("="*80)
    
    # Create new orchestrator
    orchestrator = PromptOrchestrator()
    
    # Initially, no modifiers should be loaded
    print("\nBefore accessing any modifiers:")
    print(f"  _modifiers: {orchestrator._modifiers}")
    print(f"  _weather_modifiers: {orchestrator._weather_modifiers}")
    
    # Access one modifier type
    print("\nAfter accessing lighting modifiers:")
    _ = orchestrator.modifiers
    print(f"  _modifiers: Loaded ({len(orchestrator._modifiers)} items)")
    print(f"  _weather_modifiers: {orchestrator._weather_modifiers}")
    
    # Access another modifier type
    print("\nAfter accessing weather modifiers:")
    _ = orchestrator.weather_modifiers
    print(f"  _weather_modifiers: Loaded ({len(orchestrator._weather_modifiers)} items)")
    
    print("\n✓ Lazy loading working correctly!")


def test_sample_combinations():
    """Test combining multiple modifier types"""
    orchestrator = PromptOrchestrator()
    base = "Sports car"
    
    print("\n" + "="*80)
    print("COMBINATION EXAMPLES")
    print("="*80)
    
    # Example 1: Location + Seasonal + Lighting
    result = base
    result = orchestrator.apply_location_modifier(result, 16)  # Coastal
    result = orchestrator.apply_seasonal_modifier(result, 11)  # Early autumn
    result = orchestrator.apply_lighting_modifier(result, 7)   # Golden hour
    
    print("\nExample 1: Location + Seasonal + Lighting")
    print(f"Base: {base}")
    print(f"Result: {result}")
    
    # Example 2: Weather + Mood + Camera
    result = base
    result = orchestrator.apply_weather_modifier(result, 18)  # Post-storm
    result = orchestrator.apply_mood_modifier(result, 3)      # Dramatic tension
    result = orchestrator.apply_camera_modifier(result, 7)    # Low angle
    
    print("\nExample 2: Weather + Mood + Camera")
    print(f"Base: {base}")
    print(f"Result: {result}")
    
    # Example 3: Color + Action + Composition
    result = base
    result = orchestrator.apply_color_grading_modifier(result, 17)  # Teal/orange
    result = orchestrator.apply_action_modifier(result, 1)          # In motion
    result = orchestrator.apply_composition_modifier(result, 1)      # Rule of thirds
    
    print("\nExample 3: Color + Action + Composition")
    print(f"Base: {base}")
    print(f"Result: {result}")


if __name__ == "__main__":
    print("\n" + "="*80)
    print("MODULAR ARCHITECTURE COMPREHENSIVE TEST")
    print("="*80)
    
    # Run all tests
    success = test_all_modifier_types()
    test_wet_dry_distribution()
    test_lazy_loading()
    test_sample_combinations()
    
    print("\n" + "="*80)
    if success:
        print("✓ ALL TESTS PASSED - MODULAR SYSTEM WORKING CORRECTLY")
    else:
        print("✗ SOME TESTS FAILED - CHECK OUTPUT ABOVE")
    print("="*80 + "\n")