241 lines
6.6 KiB
Python
241 lines
6.6 KiB
Python
import asyncio
|
|
import importlib.util
|
|
import sys
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
|
|
PACKAGE_ROOT = Path(__file__).parents[1]
|
|
COMFYUI_ROOT = PACKAGE_ROOT.parents[1]
|
|
sys.path.insert(0, str(COMFYUI_ROOT))
|
|
SPEC = importlib.util.spec_from_file_location("lazy_fallback_switch", PACKAGE_ROOT / "__init__.py")
|
|
MODULE = importlib.util.module_from_spec(SPEC)
|
|
SPEC.loader.exec_module(MODULE)
|
|
|
|
import nodes
|
|
from comfy_api.latest import _io
|
|
from comfy_execution.graph import DynamicPrompt, TopologicalSort
|
|
|
|
|
|
class SourceNode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required": {}}
|
|
|
|
|
|
class LazyFallbackSwitchTests(unittest.TestCase):
|
|
def test_requests_connected_selected_value(self):
|
|
result = MODULE.LazyFallbackSwitchAutogrow.check_lazy_status(
|
|
index=3,
|
|
default=None,
|
|
values={"value3": (None, "value3")},
|
|
)
|
|
|
|
self.assertEqual(result, ["value3"])
|
|
|
|
def test_does_not_request_default_when_selected_value_is_ready(self):
|
|
result = MODULE.LazyFallbackSwitchAutogrow.check_lazy_status(
|
|
index=3,
|
|
default=None,
|
|
values={"value3": ("selected", "value3")},
|
|
)
|
|
|
|
self.assertEqual(result, [])
|
|
|
|
def test_requests_default_for_unconnected_selected_value(self):
|
|
result = MODULE.LazyFallbackSwitchAutogrow.check_lazy_status(
|
|
index=3,
|
|
default=None,
|
|
values={"value2": (None, "value2")},
|
|
)
|
|
|
|
self.assertEqual(result, ["default"])
|
|
|
|
def test_does_not_request_unconnected_optional_default(self):
|
|
result = MODULE.LazyFallbackSwitchAutogrow.check_lazy_status(
|
|
index=3,
|
|
values={"value2": (None, "value2")},
|
|
)
|
|
|
|
self.assertEqual(result, [])
|
|
|
|
def test_returns_none_when_default_and_selected_value_are_unconnected(self):
|
|
output = MODULE.LazyFallbackSwitchAutogrow.execute(
|
|
index=3,
|
|
values={"value2": "other"},
|
|
)
|
|
|
|
self.assertEqual(output.args, (None,))
|
|
|
|
def test_uses_default_for_out_of_range_index(self):
|
|
output = MODULE.LazyFallbackSwitchAutogrow.execute(
|
|
index=100,
|
|
default="fallback",
|
|
values={"value99": "last"},
|
|
)
|
|
|
|
self.assertEqual(output.args, ("fallback",))
|
|
|
|
def test_returns_selected_value(self):
|
|
output = MODULE.LazyFallbackSwitchAutogrow.execute(
|
|
index=3,
|
|
default="fallback",
|
|
values={"value3": "selected"},
|
|
)
|
|
|
|
self.assertEqual(output.args, ("selected",))
|
|
|
|
def test_extension_registers_node(self):
|
|
extension = asyncio.run(MODULE.comfy_entrypoint())
|
|
nodes = asyncio.run(extension.get_node_list())
|
|
|
|
self.assertEqual(
|
|
nodes,
|
|
[MODULE.LazyFallbackSwitch, MODULE.LazyFallbackSwitchAutogrow],
|
|
)
|
|
|
|
def test_comfyui_preserves_dynamic_input_name_for_lazy_request(self):
|
|
live_inputs = {
|
|
"index": 3,
|
|
"default": ["default_node", 0],
|
|
"values.value3": ["selected_node", 0],
|
|
}
|
|
_, _, v3_data = _io.get_finalized_class_inputs(
|
|
MODULE.LazyFallbackSwitchAutogrow.INPUT_TYPES(),
|
|
live_inputs,
|
|
)
|
|
v3_data["create_dynamic_tuple"] = True
|
|
inputs = _io.build_nested_inputs(
|
|
{"index": 3, "default": None, "values.value3": None},
|
|
v3_data,
|
|
)
|
|
|
|
self.assertEqual(inputs["values"]["value3"], (None, "values.value3"))
|
|
self.assertEqual(
|
|
MODULE.LazyFallbackSwitchAutogrow.check_lazy_status(**inputs),
|
|
["values.value3"],
|
|
)
|
|
|
|
def test_autogrow_schema_starts_with_value0_through_value9(self):
|
|
inputs = MODULE.LazyFallbackSwitchAutogrow.INPUT_TYPES()
|
|
|
|
self.assertIn("default", inputs["optional"])
|
|
self.assertEqual(inputs["optional"]["values"][1]["template"]["min"], 9)
|
|
|
|
def test_finalized_value_inputs_are_optional(self):
|
|
inputs, _, _ = _io.get_finalized_class_inputs(
|
|
MODULE.LazyFallbackSwitchAutogrow.INPUT_TYPES(),
|
|
{"index": 5},
|
|
)
|
|
|
|
self.assertEqual(list(inputs["required"]), ["index"])
|
|
for index in range(10):
|
|
self.assertIn(f"values.value{index}", inputs["optional"])
|
|
|
|
def test_comfyui_scheduler_keeps_unselected_routes_lazy(self):
|
|
prompt = {
|
|
"route3": {"class_type": "SourceNode", "inputs": {}},
|
|
"route5": {"class_type": "SourceNode", "inputs": {}},
|
|
"switch": {
|
|
"class_type": "SparknightLazyFallbackSwitch",
|
|
"inputs": {
|
|
"index": 3,
|
|
"values.value3": ["route3", 0],
|
|
"values.value5": ["route5", 0],
|
|
},
|
|
},
|
|
}
|
|
|
|
with patch.dict(
|
|
nodes.NODE_CLASS_MAPPINGS,
|
|
{
|
|
"SourceNode": SourceNode,
|
|
"SparknightLazyFallbackSwitch": MODULE.LazyFallbackSwitchAutogrow,
|
|
},
|
|
):
|
|
sort = TopologicalSort(DynamicPrompt(prompt))
|
|
sort.add_node("switch")
|
|
self.assertEqual(set(sort.pendingNodes), {"switch"})
|
|
|
|
sort.make_input_strong_link("switch", "values.value3")
|
|
self.assertEqual(set(sort.pendingNodes), {"switch", "route3"})
|
|
|
|
def test_static_schema_has_only_value0_through_value9(self):
|
|
inputs = MODULE.LazyFallbackSwitch.INPUT_TYPES()
|
|
|
|
self.assertEqual(list(inputs["required"]), ["index"])
|
|
self.assertIn("default", inputs["optional"])
|
|
for index in range(10):
|
|
self.assertIn(f"value{index}", inputs["optional"])
|
|
self.assertTrue(inputs["optional"][f"value{index}"][1]["lazy"])
|
|
self.assertNotIn("value10", inputs["optional"])
|
|
|
|
def test_static_switch_requests_only_selected_route(self):
|
|
result = MODULE.LazyFallbackSwitch.check_lazy_status(
|
|
index=3,
|
|
default=None,
|
|
value3=None,
|
|
value5="unselected",
|
|
)
|
|
|
|
self.assertEqual(result, ["value3"])
|
|
|
|
def test_comfyui_scheduler_keeps_unselected_static_routes_lazy(self):
|
|
prompt = {
|
|
"route3": {"class_type": "SourceNode", "inputs": {}},
|
|
"route5": {"class_type": "SourceNode", "inputs": {}},
|
|
"switch": {
|
|
"class_type": "SparknightLazyFallbackSwitchStatic",
|
|
"inputs": {
|
|
"index": 3,
|
|
"value3": ["route3", 0],
|
|
"value5": ["route5", 0],
|
|
},
|
|
},
|
|
}
|
|
|
|
with patch.dict(
|
|
nodes.NODE_CLASS_MAPPINGS,
|
|
{
|
|
"SourceNode": SourceNode,
|
|
"SparknightLazyFallbackSwitchStatic": MODULE.LazyFallbackSwitch,
|
|
},
|
|
):
|
|
sort = TopologicalSort(DynamicPrompt(prompt))
|
|
sort.add_node("switch")
|
|
self.assertEqual(set(sort.pendingNodes), {"switch"})
|
|
|
|
sort.make_input_strong_link("switch", "value3")
|
|
self.assertEqual(set(sort.pendingNodes), {"switch", "route3"})
|
|
|
|
def test_static_switch_returns_selected_route(self):
|
|
output = MODULE.LazyFallbackSwitch.execute(
|
|
index=3,
|
|
default="fallback",
|
|
value3="selected",
|
|
value5="unselected",
|
|
)
|
|
|
|
self.assertEqual(output.args, ("selected",))
|
|
|
|
def test_static_switch_uses_default_for_missing_route(self):
|
|
output = MODULE.LazyFallbackSwitch.execute(
|
|
index=3,
|
|
default="fallback",
|
|
value5="unselected",
|
|
)
|
|
|
|
self.assertEqual(output.args, ("fallback",))
|
|
|
|
def test_display_names_distinguish_variants(self):
|
|
self.assertEqual(MODULE.LazyFallbackSwitch.define_schema().display_name, "Lazy Fallback Switch")
|
|
self.assertEqual(
|
|
MODULE.LazyFallbackSwitchAutogrow.define_schema().display_name,
|
|
"Lazy Fallback Switch (Autogrow)",
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|