325 lines
8.2 KiB
Python
325 lines
8.2 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Test Suite for ComfyUI Prompt Orchestration Nodes
|
|
Validates that all nodes work correctly before deployment
|
|
"""
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
# Add current directory to path
|
|
sys.path.insert(0, str(Path(__file__).parent))
|
|
|
|
from prompt_orchestration_nodes import (
|
|
PromptLightingModifierNode,
|
|
PromptBatchVariationNode,
|
|
PromptStyleTransformNode,
|
|
PromptSmartVariationNode,
|
|
PromptPresetCollectionNode,
|
|
PromptModifierInfoNode,
|
|
PromptCombinerNode,
|
|
PromptRandomModifierNode,
|
|
PromptMultiStyleNode
|
|
)
|
|
|
|
|
|
def test_lighting_modifier():
|
|
"""Test single lighting modifier node"""
|
|
print("\n" + "="*80)
|
|
print("TEST 1: Lighting Modifier Node")
|
|
print("="*80)
|
|
|
|
node = PromptLightingModifierNode()
|
|
|
|
# Test with golden hour
|
|
result = node.apply_modifier(
|
|
text="Sports car on highway",
|
|
modifier_id=7,
|
|
preserve_base=True
|
|
)
|
|
|
|
print(f"Input: 'Sports car on highway'")
|
|
print(f"Modifier: 7 (golden hour)")
|
|
print(f"Output: {result[0]}")
|
|
|
|
# Verify output contains expected elements
|
|
assert "Sports car on highway" in result[0], "Base text missing"
|
|
assert "golden hour" in result[0], "Modifier text missing"
|
|
|
|
print("✅ PASSED")
|
|
return True
|
|
|
|
|
|
def test_batch_variation():
|
|
"""Test batch variation generator node"""
|
|
print("\n" + "="*80)
|
|
print("TEST 2: Batch Variation Generator")
|
|
print("="*80)
|
|
|
|
node = PromptBatchVariationNode()
|
|
|
|
# Test with custom IDs
|
|
results = node.generate_variations(
|
|
text="Modern building",
|
|
category="custom",
|
|
custom_ids="1,7,18"
|
|
)
|
|
|
|
print(f"Input: 'Modern building'")
|
|
print(f"Category: custom")
|
|
print(f"Custom IDs: 1,7,18")
|
|
print(f"Generated {len(results)} outputs")
|
|
|
|
# Check first few variations
|
|
for i in range(min(3, len(results))):
|
|
print(f" var_{i+1}: {results[i][0][:80]}...")
|
|
|
|
assert len(results) == 10, "Should return 10 outputs"
|
|
|
|
print("✅ PASSED")
|
|
return True
|
|
|
|
|
|
def test_style_transform():
|
|
"""Test style transformation node"""
|
|
print("\n" + "="*80)
|
|
print("TEST 3: Style Transform Node")
|
|
print("="*80)
|
|
|
|
node = PromptStyleTransformNode()
|
|
|
|
result = node.apply_style(
|
|
text="Vintage car in garage",
|
|
style="cyberpunk",
|
|
maintain_composition=True
|
|
)
|
|
|
|
print(f"Input: 'Vintage car in garage'")
|
|
print(f"Style: cyberpunk")
|
|
print(f"Output: {result[0][:100]}...")
|
|
|
|
assert "cyberpunk" in result[0].lower(), "Style not applied"
|
|
|
|
print("✅ PASSED")
|
|
return True
|
|
|
|
|
|
def test_smart_variation():
|
|
"""Test smart variation with natural language"""
|
|
print("\n" + "="*80)
|
|
print("TEST 4: Smart Variation Node")
|
|
print("="*80)
|
|
|
|
node = PromptSmartVariationNode()
|
|
|
|
result = node.smart_vary(
|
|
text="Car on street",
|
|
intent="make it golden hour"
|
|
)
|
|
|
|
print(f"Input: 'Car on street'")
|
|
print(f"Intent: 'make it golden hour'")
|
|
print(f"Output: {result[0]}")
|
|
|
|
assert "golden hour" in result[0].lower() or "warm" in result[0].lower(), "Intent not recognized"
|
|
|
|
print("✅ PASSED")
|
|
return True
|
|
|
|
|
|
def test_preset_collection():
|
|
"""Test preset collection node"""
|
|
print("\n" + "="*80)
|
|
print("TEST 5: Preset Collection Node")
|
|
print("="*80)
|
|
|
|
node = PromptPresetCollectionNode()
|
|
|
|
results = node.apply_preset(
|
|
text="Product on pedestal",
|
|
preset="commercial"
|
|
)
|
|
|
|
print(f"Input: 'Product on pedestal'")
|
|
print(f"Preset: commercial")
|
|
print(f"Generated {len(results)} outputs")
|
|
|
|
for i in range(min(4, len(results))):
|
|
print(f" var_{i+1}: {results[i][0][:60]}...")
|
|
|
|
assert len(results) == 8, "Should return 8 outputs"
|
|
|
|
print("✅ PASSED")
|
|
return True
|
|
|
|
|
|
def test_modifier_info():
|
|
"""Test modifier info node"""
|
|
print("\n" + "="*80)
|
|
print("TEST 6: Modifier Info Node")
|
|
print("="*80)
|
|
|
|
node = PromptModifierInfoNode()
|
|
|
|
desc, category, tags = node.get_info(modifier_id=7)
|
|
|
|
print(f"Modifier ID: 7")
|
|
print(f"Description: {desc}")
|
|
print(f"Category: {category}")
|
|
print(f"Tags: {tags}")
|
|
|
|
assert "golden hour" in desc.lower(), "Wrong modifier info"
|
|
assert category in ["dry", "wet"], "Invalid category"
|
|
|
|
print("✅ PASSED")
|
|
return True
|
|
|
|
|
|
def test_combiner():
|
|
"""Test prompt combiner node"""
|
|
print("\n" + "="*80)
|
|
print("TEST 7: Prompt Combiner Node")
|
|
print("="*80)
|
|
|
|
node = PromptCombinerNode()
|
|
|
|
result = node.combine(
|
|
text_1="Sports car",
|
|
text_2="on mountain road",
|
|
text_3="golden hour",
|
|
separator=", "
|
|
)
|
|
|
|
print(f"Input 1: 'Sports car'")
|
|
print(f"Input 2: 'on mountain road'")
|
|
print(f"Input 3: 'golden hour'")
|
|
print(f"Output: {result[0]}")
|
|
|
|
assert "Sports car" in result[0], "Missing text_1"
|
|
assert "mountain road" in result[0], "Missing text_2"
|
|
assert "golden hour" in result[0], "Missing text_3"
|
|
|
|
print("✅ PASSED")
|
|
return True
|
|
|
|
|
|
def test_random_modifier():
|
|
"""Test random modifier node"""
|
|
print("\n" + "="*80)
|
|
print("TEST 8: Random Modifier Node")
|
|
print("="*80)
|
|
|
|
node = PromptRandomModifierNode()
|
|
|
|
result, modifier_id = node.apply_random(
|
|
text="Architecture scene",
|
|
category="dry",
|
|
seed=42
|
|
)
|
|
|
|
print(f"Input: 'Architecture scene'")
|
|
print(f"Category: dry")
|
|
print(f"Seed: 42")
|
|
print(f"Selected Modifier: {modifier_id}")
|
|
print(f"Output: {result[:80]}...")
|
|
|
|
assert 1 <= modifier_id <= 16, "Invalid modifier ID for DRY category"
|
|
|
|
# Test reproducibility
|
|
result2, modifier_id2 = node.apply_random(
|
|
text="Architecture scene",
|
|
category="dry",
|
|
seed=42
|
|
)
|
|
|
|
assert modifier_id == modifier_id2, "Same seed should give same modifier"
|
|
|
|
print("✅ PASSED (reproducible)")
|
|
return True
|
|
|
|
|
|
def test_multi_style():
|
|
"""Test multi-style layer node"""
|
|
print("\n" + "="*80)
|
|
print("TEST 9: Multi-Style Layer Node")
|
|
print("="*80)
|
|
|
|
node = PromptMultiStyleNode()
|
|
|
|
result = node.apply_multi_style(
|
|
text="Urban scene",
|
|
style_1="cyberpunk",
|
|
style_2="noir",
|
|
style_3="none"
|
|
)
|
|
|
|
print(f"Input: 'Urban scene'")
|
|
print(f"Style 1: cyberpunk")
|
|
print(f"Style 2: noir")
|
|
print(f"Output: {result[0][:100]}...")
|
|
|
|
# Should contain elements of both styles
|
|
assert len(result[0]) > len("Urban scene"), "Styles not applied"
|
|
|
|
print("✅ PASSED")
|
|
return True
|
|
|
|
|
|
def run_all_tests():
|
|
"""Run all tests and report results"""
|
|
print("\n" + "█"*80)
|
|
print("COMFYUI PROMPT ORCHESTRATION NODES - TEST SUITE")
|
|
print("█"*80)
|
|
|
|
tests = [
|
|
("Lighting Modifier", test_lighting_modifier),
|
|
("Batch Variation", test_batch_variation),
|
|
("Style Transform", test_style_transform),
|
|
("Smart Variation", test_smart_variation),
|
|
("Preset Collection", test_preset_collection),
|
|
("Modifier Info", test_modifier_info),
|
|
("Prompt Combiner", test_combiner),
|
|
("Random Modifier", test_random_modifier),
|
|
("Multi-Style Layer", test_multi_style)
|
|
]
|
|
|
|
results = []
|
|
|
|
for test_name, test_func in tests:
|
|
try:
|
|
success = test_func()
|
|
results.append((test_name, True, None))
|
|
except Exception as e:
|
|
print(f"\n❌ FAILED: {e}")
|
|
results.append((test_name, False, str(e)))
|
|
|
|
# Summary
|
|
print("\n" + "█"*80)
|
|
print("TEST SUMMARY")
|
|
print("█"*80)
|
|
|
|
passed = sum(1 for _, success, _ in results if success)
|
|
total = len(results)
|
|
|
|
for test_name, success, error in results:
|
|
status = "✅ PASS" if success else "❌ FAIL"
|
|
print(f"{status} - {test_name}")
|
|
if error:
|
|
print(f" Error: {error}")
|
|
|
|
print("\n" + "="*80)
|
|
print(f"Results: {passed}/{total} tests passed")
|
|
print("="*80)
|
|
|
|
if passed == total:
|
|
print("\n🎉 ALL TESTS PASSED! Nodes are ready for ComfyUI.")
|
|
return True
|
|
else:
|
|
print(f"\n⚠️ {total - passed} test(s) failed. Please review errors above.")
|
|
return False
|
|
|
|
|
|
if __name__ == "__main__":
|
|
success = run_all_tests()
|
|
sys.exit(0 if success else 1)
|