159 lines
5.7 KiB
Python
159 lines
5.7 KiB
Python
"""
|
|
Tests for R12 WorkflowPruner service.
|
|
"""
|
|
|
|
import unittest
|
|
import sys
|
|
import os
|
|
from typing import Dict, Any
|
|
|
|
# --- PATH SETUP ---
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
project_root = os.path.abspath(os.path.join(current_dir, ".."))
|
|
if project_root not in sys.path:
|
|
sys.path.insert(0, project_root)
|
|
|
|
from services.workflow_pruner import WorkflowPruner, PruneConfig
|
|
|
|
class TestWorkflowPruner(unittest.TestCase):
|
|
|
|
def setUp(self):
|
|
self.pruner = WorkflowPruner()
|
|
self.config = PruneConfig(
|
|
max_depth=3,
|
|
max_nodes=10,
|
|
fallback_recent_nodes_count=2
|
|
)
|
|
|
|
def test_prune_basic_upstream(self):
|
|
"""Test basic upstream pruning."""
|
|
# Simple chain: 1 -> 2 -> 3(target)
|
|
workflow = {
|
|
"1": {"class_type": "Loader", "inputs": {}},
|
|
"2": {"class_type": "Sampler", "inputs": {"model": ["1", 0]}},
|
|
"3": {"class_type": "Save", "inputs": {"images": ["2", 0]}},
|
|
"4": {"class_type": "Unrelated", "inputs": {}}
|
|
}
|
|
|
|
result = self.pruner.prune(workflow, "3", self.config)
|
|
|
|
self.assertEqual(result.mode, "upstream_trace")
|
|
self.assertIn("3", result.kept_node_ids)
|
|
self.assertIn("2", result.kept_node_ids)
|
|
self.assertIn("1", result.kept_node_ids)
|
|
self.assertNotIn("4", result.kept_node_ids) # Unrelated should be pruned
|
|
self.assertEqual(result.nodes_used, 3)
|
|
|
|
def test_prune_depth_limit(self):
|
|
"""Test max depth limit."""
|
|
# Chain: 1 -> 2 -> 3 -> 4(target)
|
|
# Depth limit 2 means: 4(0), 3(1), 2(2). Node 1 is depth 3.
|
|
# Actually in BFS: 4 is depth 0.
|
|
# Neighbors of 4 (node 3) is depth 1.
|
|
# Neighbors of 3 (node 2) is depth 2.
|
|
# Neighbors of 2 (node 1) is depth 3.
|
|
# If max_depth=2, we keep depths 0, 1, 2. (So 4, 3, 2). Node 1 should be dropped.
|
|
|
|
config = PruneConfig(max_depth=2, max_nodes=10)
|
|
|
|
workflow = {
|
|
"1": {"inputs": {}},
|
|
"2": {"inputs": {"in": ["1", 0]}},
|
|
"3": {"inputs": {"in": ["2", 0]}},
|
|
"4": {"inputs": {"in": ["3", 0]}}
|
|
}
|
|
|
|
result = self.pruner.prune(workflow, "4", config)
|
|
|
|
self.assertIn("4", result.kept_node_ids) # depth 0
|
|
self.assertIn("3", result.kept_node_ids) # depth 1
|
|
self.assertIn("2", result.kept_node_ids) # depth 2
|
|
self.assertNotIn("1", result.kept_node_ids) # depth 3 >= max_depth 2?
|
|
# Wait, if logic is `if current_depth >= config.max_depth: continue`
|
|
# Depth 0 (target) -> expand depth 1.
|
|
# Depth 1 -> expand depth 2.
|
|
# Depth 2 -> expands depth 3?
|
|
# If max_depth is 2.
|
|
# Pop(4, 0). 0 < 2. Add 3 (1).
|
|
# Pop(3, 1). 1 < 2. Add 2 (2).
|
|
# Pop(2, 2). 2 >= 2. Continue. (Node 1 not added).
|
|
# So yes, Node 1 should be excluded.
|
|
|
|
self.assertEqual(len(result.kept_node_ids), 3)
|
|
|
|
def test_prune_node_limit(self):
|
|
"""Test max nodes limit."""
|
|
# 1 -> 2 -> 3(target). Max nodes = 2.
|
|
config = PruneConfig(max_depth=5, max_nodes=2)
|
|
|
|
workflow = {
|
|
"1": {"inputs": {}},
|
|
"2": {"inputs": {"in": ["1", 0]}},
|
|
"3": {"inputs": {"in": ["2", 0]}}
|
|
}
|
|
|
|
result = self.pruner.prune(workflow, "3", config)
|
|
|
|
self.assertIn("3", result.kept_node_ids)
|
|
# Should keep one more. Usually BFS order determines likelihood.
|
|
# 3 depends on 2. So 2 is added.
|
|
# 2 depends on 1. But limit reached.
|
|
|
|
self.assertEqual(len(result.kept_node_ids), 2)
|
|
self.assertIn("3", result.kept_node_ids)
|
|
self.assertIn("2", result.kept_node_ids)
|
|
self.assertNotIn("1", result.kept_node_ids)
|
|
|
|
def test_fallback_missing_target(self):
|
|
"""Test fallback when target node missing."""
|
|
workflow = {
|
|
"1": {"inputs": {}},
|
|
"2": {"inputs": {}},
|
|
"3": {"inputs": {}}
|
|
}
|
|
|
|
# Target "99" missing
|
|
result = self.pruner.prune(workflow, "99", self.config)
|
|
|
|
self.assertEqual(result.mode, "fallback_recent_nodes")
|
|
# specific config set fallback count to 2
|
|
self.assertEqual(len(result.kept_node_ids), 2)
|
|
# Should be last 2 by ID sort: 2, 3
|
|
self.assertIn("3", result.kept_node_ids)
|
|
self.assertIn("2", result.kept_node_ids)
|
|
self.assertNotIn("1", result.kept_node_ids)
|
|
|
|
def test_deterministic_traversal(self):
|
|
"""Test that traversal is deterministic (order of inputs)."""
|
|
# Node 3 depends on 2 and 1.
|
|
# If traversal isn't sorted, we might visit 1 or 2 first.
|
|
# With sorting, we should always visit 1 then 2 (if keys are "a", "b") or based on key names.
|
|
|
|
workflow = {
|
|
"1": {"inputs": {}},
|
|
"2": {"inputs": {}},
|
|
"3": {"inputs": {
|
|
"b_input": ["2", 0],
|
|
"a_input": ["1", 0]
|
|
}}
|
|
}
|
|
|
|
# Input keys: "a_input", "b_input". Sorted: a, b.
|
|
# Should visit 1 ("a_input") then 2 ("b_input").
|
|
|
|
# To test this, we set max_nodes=2 (target + 1 neighbor).
|
|
# Target is 3.
|
|
# If "a" comes first, 1 is visited. Limit reached. 2 excluded.
|
|
# If "b" comes first, 2 is visited. Limit reached. 1 excluded.
|
|
|
|
config = PruneConfig(max_nodes=2)
|
|
result = self.pruner.prune(workflow, "3", config)
|
|
|
|
self.assertIn("3", result.kept_node_ids)
|
|
self.assertIn("1", result.kept_node_ids) # "a_input" -> node 1
|
|
self.assertNotIn("2", result.kept_node_ids) # "b_input" -> node 2 (dropped)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|