eniewold-ComfyUI-Subworkflow/tests/test_workflow_node.py

import importlib.util
import logging
import pathlib
import sys
import types
import unittest
from types import SimpleNamespace


REPO_ROOT = pathlib.Path(__file__).resolve().parents[1]


def _load_workflow_node_module():
    package_name = "swf_isolated_test"
    package = types.ModuleType(package_name)
    package.__path__ = [str(REPO_ROOT)]
    sys.modules[package_name] = package

    latest = types.ModuleType("comfy_api.latest")
    latest.io = SimpleNamespace(
        ComfyNode=object,
        NodeOutput=lambda *args, **kwargs: (args, kwargs),
        Hidden=SimpleNamespace(unique_id="UNIQUE_ID", dynprompt="DYNPROMPT"),
    )
    comfy_api = types.ModuleType("comfy_api")
    comfy_api.latest = latest
    sys.modules["comfy_api"] = comfy_api
    sys.modules["comfy_api.latest"] = latest

    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

    workflow_utils = types.ModuleType(f"{package_name}.workflow_utils")
    for name in (
        "list_workflow_files",
        "load_workflow_file",
        "load_workflow_url",
        "build_expansion",
        "build_modifier_source_expansion",
        "apply_control_after_generate",
        "validate_workflow_nodes_installed",
        "get_workflow_interface",
    ):
        setattr(workflow_utils, name, lambda *args, **kwargs: None)
    sys.modules[workflow_utils.__name__] = workflow_utils

    module_name = f"{package_name}.workflow_node"
    spec = importlib.util.spec_from_file_location(
        module_name,
        REPO_ROOT / "workflow_node.py",
    )
    module = importlib.util.module_from_spec(spec)
    sys.modules[module_name] = module
    spec.loader.exec_module(module)
    return module


workflow_node = _load_workflow_node_module()


class FakePrompt:
    def __init__(self, nodes=None):
        self.nodes = nodes or {}

    def has_node(self, node_id):
        return node_id in self.nodes

    def get_node(self, node_id):
        return self.nodes[node_id]


class ModifierSourceTargetTests(unittest.TestCase):
    def test_resolves_file_and_url_subworkflow_targets(self):
        prompt = FakePrompt({
            "12": {
                "class_type": "SWF_Subworkflow",
                "inputs": {
                    "workflow": "inner.json",
                    "reload_each_execution": False,
                },
            },
            "13": {
                "class_type": "SWF_SubworkflowFromURL",
                "inputs": {
                    "url": "https://example.test/workflow.json",
                    "verify_ssl": False,
                    "reload_each_execution": True,
                },
            },
        })
        workflow_node.SubworkflowModifierSource.hidden = SimpleNamespace(
            dynprompt=prompt
        )

        file_target = workflow_node.SubworkflowModifierSource._resolve_target(
            "12 | Local | inner.json"
        )
        url_target = workflow_node.SubworkflowModifierSource._resolve_target(
            "13 | Remote | https://example.test/workflow.json"
        )

        self.assertIs(file_target[1], workflow_node.Subworkflow)
        self.assertEqual(file_target[2], {"workflow": "inner.json"})
        self.assertIs(file_target[3], False)
        self.assertIs(url_target[1], workflow_node.SubworkflowFromURL)
        self.assertEqual(url_target[2]["verify_ssl"], False)
        self.assertIs(url_target[3], True)

    def test_rejects_non_subworkflow_target(self):
        prompt = FakePrompt({
            "20": {"class_type": "KSampler", "inputs": {}},
        })
        workflow_node.SubworkflowModifierSource.hidden = SimpleNamespace(
            dynprompt=prompt
        )
        with self.assertRaisesRegex(ValueError, "not a Subworkflow"):
            workflow_node.SubworkflowModifierSource._resolve_target("20 | Sampler")


class WorkflowDataCacheTests(unittest.TestCase):
    def setUp(self):
        workflow_node.BaseSubworkflow._loaded_workflows.clear()
        workflow_node.BaseSubworkflow._loaded_workflow_instances.clear()
        workflow_node.BaseSubworkflow._prompt_workflows.clear()
        self.loads = []
        self.original_loader = workflow_node.Subworkflow.__dict__["_load_source"]
        loads = self.loads
        workflow_node.Subworkflow._load_source = classmethod(
            lambda cls, workflow, **kwargs: (
                loads.append(workflow) or {"load": len(loads)}
            )
        )

    def tearDown(self):
        workflow_node.Subworkflow._load_source = self.original_loader

    def test_reload_loads_once_per_prompt(self):
        prompt = FakePrompt()
        first = workflow_node.Subworkflow._get_workflow_data(
            True,
            instance_id="12",
            dynprompt=prompt,
            workflow="inner.json",
        )
        second = workflow_node.Subworkflow._get_workflow_data(
            True,
            instance_id="12",
            dynprompt=prompt,
            workflow="inner.json",
        )
        self.assertIs(first, second)
        self.assertEqual(len(self.loads), 1)

        workflow_node.Subworkflow._get_workflow_data(
            True,
            instance_id="12",
            dynprompt=FakePrompt(),
            workflow="inner.json",
        )
        self.assertEqual(len(self.loads), 2)

    def test_keep_loaded_is_isolated_by_target_node(self):
        kept = workflow_node.Subworkflow._get_workflow_data(
            False,
            instance_id="12",
            dynprompt=FakePrompt(),
            workflow="inner.json",
        )
        kept_again = workflow_node.Subworkflow._get_workflow_data(
            False,
            instance_id="12",
            dynprompt=FakePrompt(),
            workflow="inner.json",
        )
        self.assertIs(kept, kept_again)
        self.assertEqual(len(self.loads), 1)

        workflow_node.Subworkflow._get_workflow_data(
            False,
            instance_id="99",
            dynprompt=FakePrompt(),
            workflow="inner.json",
        )
        self.assertEqual(len(self.loads), 2)


if __name__ == "__main__":
    unittest.main()