76 lines
2.1 KiB
Python
76 lines
2.1 KiB
Python
import importlib.util
|
|
import logging
|
|
import pathlib
|
|
import sys
|
|
import types
|
|
import unittest
|
|
|
|
|
|
REPO_ROOT = pathlib.Path(__file__).resolve().parents[1]
|
|
|
|
|
|
def _load_workflow_utils_module():
|
|
package_name = "swf_utils_isolated_test"
|
|
package = types.ModuleType(package_name)
|
|
package.__path__ = [str(REPO_ROOT)]
|
|
sys.modules[package_name] = package
|
|
|
|
graph_utils = types.ModuleType("comfy_execution.graph_utils")
|
|
graph_utils.GraphBuilder = object
|
|
comfy_execution = types.ModuleType("comfy_execution")
|
|
comfy_execution.graph_utils = graph_utils
|
|
sys.modules["comfy_execution"] = comfy_execution
|
|
sys.modules["comfy_execution.graph_utils"] = graph_utils
|
|
|
|
debug_utils = types.ModuleType(f"{package_name}.debug_utils")
|
|
debug_utils.configure_logger = lambda name: logging.getLogger(name)
|
|
sys.modules[debug_utils.__name__] = debug_utils
|
|
|
|
module_name = f"{package_name}.workflow_utils"
|
|
spec = importlib.util.spec_from_file_location(
|
|
module_name,
|
|
REPO_ROOT / "workflow_utils.py",
|
|
)
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules[module_name] = module
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
workflow_utils = _load_workflow_utils_module()
|
|
|
|
|
|
class FakeNode:
|
|
def __init__(self, **inputs):
|
|
self.inputs = inputs
|
|
|
|
|
|
class FakeGraph:
|
|
def __init__(self, nodes):
|
|
self.nodes = nodes
|
|
|
|
|
|
class ModifierSourcePruningTests(unittest.TestCase):
|
|
def test_keeps_only_selected_output_ancestors(self):
|
|
graph = FakeGraph({
|
|
"source": FakeNode(value=1),
|
|
"modifier_value": FakeNode(value=["source", 0]),
|
|
"sampler": FakeNode(model=["modifier_value", 0]),
|
|
"save": FakeNode(images=["sampler", 0]),
|
|
"selected": FakeNode(nested={"value": ["modifier_value", 0]}),
|
|
})
|
|
|
|
workflow_utils._prune_graph_to_output_refs(
|
|
graph,
|
|
[["selected", 0]],
|
|
)
|
|
|
|
self.assertEqual(
|
|
set(graph.nodes),
|
|
{"source", "modifier_value", "selected"},
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|