From c49f3832afd028ddecd10ddae2b9b2e28c64bd74 Mon Sep 17 00:00:00 2001 From: WildAi <2853742+wildminder@users.noreply.github.com> Date: Fri, 23 Oct 2026 20:55:15 +0300 Subject: [PATCH] plan(P4.1): migrate the model loader node to the V3 schema --- conftest.py | 120 +++++++-------- nodes/keep_model_loader.py | 102 +++++++++---- tests/test_nodes_execute.py | 139 ++++++++++++++++- tests/test_nodes_layout.py | 290 ++++++++++++++++++++++++++++-------- 4 files changed, 501 insertions(+), 150 deletions(-) diff --git a/conftest.py b/conftest.py index a1dc8f2..17b6dc2 100644 --- a/conftest.py +++ b/conftest.py @@ -193,6 +193,11 @@ class _MockSchema: class _MockInput: """Mock of ``comfy_api.latest.io.Input``.""" + #: The socket type of the input. The real API reads it from the parent + #: type (``self.Parent.io_type``); the mock carries it on the class, so + #: ``get_io_type()`` answers the same question. + io_type: str | None = None + def __init__( self, id: str, @@ -222,11 +227,16 @@ class _MockInput: "optional": self.optional, "tooltip": self.tooltip, "lazy": self.lazy, - "rawLink": self.rawLink, + "rawLink": self.raw_link, "advanced": self.advanced, **self.extra_dict, } + def get_io_type(self) -> str | None: + """The socket type of this input (the real API has the same method).""" + return self.io_type + + class _MockWidgetInput(_MockInput): """Mock of the widget inputs (``io.Int.Input``, ``io.Combo.Input``, ...). @@ -268,6 +278,9 @@ class _MockWidgetInput(_MockInput): class _MockOutput: """Mock of ``comfy_api.latest.io.Output``.""" + #: The socket type of the output (see :class:`_MockInput`). + io_type: str | None = None + def __init__( self, display_name: str | None = None, @@ -279,6 +292,10 @@ class _MockOutput: for key, value in extra.items(): setattr(self, key, value) + def get_io_type(self) -> str | None: + """The socket type of this output (the real API has the same method).""" + return self.io_type + class _MockNodeOutput: """Mock of ``comfy_api.latest.io.NodeOutput``. @@ -355,6 +372,25 @@ class _MockTypeIO: type_name: str = "" +def _mock_socket_type(io_type: str, widget: bool = False): + """A mock socket type with its own ``Input`` / ``Output`` classes. + + The real API builds a socket type per call (``io.Custom("X")``) and the + input / output classes answer their socket type through + ``get_io_type()``. The mock does the same, so a schema snapshot can ask + an input for its type (the V3 schema is the only place the socket types + live after the V1 ``INPUT_TYPES`` are gone). + """ + base_input = _MockWidgetInput if widget else _MockInput + input_cls = type(f"{io_type}_Input", (base_input,), {"io_type": io_type}) + output_cls = type(f"{io_type}_Output", (_MockOutput,), {"io_type": io_type}) + return type( + io_type, + (_MockTypeIO,), + {"type_name": io_type, "Input": input_cls, "Output": output_cls}, + ) + + class _MockIO: """Mock of ``comfy_api.latest.io`` (the *sibling* of ``ui``).""" @@ -364,59 +400,23 @@ class _MockIO: NodeOutput = _MockNodeOutput ComfyNode = _MockComfyNode - class Boolean(_MockTypeIO): - type_name = "BOOLEAN" - Input = _MockWidgetInput + Boolean = _mock_socket_type("BOOLEAN", widget=True) + Int = _mock_socket_type("INT", widget=True) + Float = _mock_socket_type("FLOAT", widget=True) + String = _mock_socket_type("STRING", widget=True) + Combo = _mock_socket_type("COMBO", widget=True) + MultiCombo = _mock_socket_type("MULTICOMBO", widget=True) - class Int(_MockTypeIO): - type_name = "INT" - Input = _MockWidgetInput - - class Float(_MockTypeIO): - type_name = "FLOAT" - Input = _MockWidgetInput - - class String(_MockTypeIO): - type_name = "STRING" - Input = _MockWidgetInput - - class Combo(_MockTypeIO): - type_name = "COMBO" - Input = _MockWidgetInput - - class MultiCombo(_MockTypeIO): - type_name = "MULTICOMBO" - Input = _MockWidgetInput - - class Image(_MockTypeIO): - type_name = "IMAGE" - - class Mask(_MockTypeIO): - type_name = "MASK" - - class Latent(_MockTypeIO): - type_name = "LATENT" - - class Conditioning(_MockTypeIO): - type_name = "CONDITIONING" - - class Model(_MockTypeIO): - type_name = "MODEL" - - class ModelPatch(_MockTypeIO): - type_name = "MODEL_PATCH" - - class Vae(_MockTypeIO): - type_name = "VAE" - - class UpscaleModel(_MockTypeIO): - type_name = "UPSCALE_MODEL" - - class Audio(_MockTypeIO): - type_name = "AUDIO" - - class Video(_MockTypeIO): - type_name = "VIDEO" + Image = _mock_socket_type("IMAGE") + Mask = _mock_socket_type("MASK") + Latent = _mock_socket_type("LATENT") + Conditioning = _mock_socket_type("CONDITIONING") + Model = _mock_socket_type("MODEL") + ModelPatch = _mock_socket_type("MODEL_PATCH") + Vae = _mock_socket_type("VAE") + UpscaleModel = _mock_socket_type("UPSCALE_MODEL") + Audio = _mock_socket_type("AUDIO") + Video = _mock_socket_type("VIDEO") class Hidden: unique_id = "unique_id" @@ -425,11 +425,7 @@ class _MockIO: @staticmethod def Custom(io_type: str): """Mock of ``io.Custom("TYPE")`` - a socket type by name.""" - return type( - f"Custom_{io_type}", - (_MockTypeIO,), - {"type_name": io_type}, - ) + return _mock_socket_type(io_type) class _MockUI: @@ -514,12 +510,16 @@ class _MockComfyAPISync: class _MockModelPatcher: """Mock of ``comfy.model_patcher.ModelPatcher``.""" - def __init__(self, model=None, *args, **kwargs) -> None: + def __init__(self, model=None, load_device=None, offload_device=None, *args, **kwargs) -> None: self.model = model if model is not None else types.SimpleNamespace() if not hasattr(self.model, "diffusion_model"): self.model.diffusion_model = types.SimpleNamespace() self._object_patches: dict = {} self._unet_wrapper = None + # The real patcher knows the devices it moves the model between + # (``ModelPatcher(model, load_device, offload_device)``). + self.load_device = load_device + self.offload_device = offload_device def clone(self) -> "_MockModelPatcher": new = _MockModelPatcher() @@ -527,6 +527,8 @@ class _MockModelPatcher: new.model.diffusion_model = copy.copy(self.model.diffusion_model) new._object_patches = dict(self._object_patches) new._unet_wrapper = self._unet_wrapper + new.load_device = self.load_device + new.offload_device = self.offload_device return new def add_object_patch(self, path: str, obj) -> None: diff --git a/nodes/keep_model_loader.py b/nodes/keep_model_loader.py index f6a6506..e44c06c 100644 --- a/nodes/keep_model_loader.py +++ b/nodes/keep_model_loader.py @@ -1,14 +1,17 @@ -"""KEEP model loader node (V1 schema - P3.1). +"""KEEP model loader node (V3 schema - P4.1). -The class body is the V1 class of the former root ``nodes.py``, moved -**verbatim**: the node id, the display name, the category, the widget names, -their defaults and their order are compatibility guarantees of the pack. +The node id (``KEEP_ModelLoader``), the display name ("Load KEEP Models"), +the category (``ComfyUI-KEEP``) and every widget (name, default, min/max/step +and order) are compatibility guarantees of the pack - the V3 schema below +carries exactly the contract of the former V1 ``INPUT_TYPES``. The engines are imported with the DyPE dual-mode import, so the node module works both as part of the node package (how ComfyUI loads it) and as a flat module (the repository root on ``sys.path``). """ +from comfy_api.latest import io + try: # the package import (ComfyUI loads the node package) from ..src.keep_model_configs import KEEP_MODEL_CONFIGS from ..src.keep_model_loader import KEEPModelLoader @@ -18,41 +21,90 @@ except ImportError: # the flat import (the repository root on sys.path) GLOBAL_KEEP_MODEL_LOADER = None +#: The face detectors of the vendored facelib helper (the V1 combo). +DETECTION_MODELS = [ + "retinaface_resnet50", + "retinaface_mobile0.25", + "YOLOv5l", + "YOLOv5n", +] + def get_keep_model_loader(): + """The global loader singleton (the cache key of the pack is unchanged).""" global GLOBAL_KEEP_MODEL_LOADER if GLOBAL_KEEP_MODEL_LOADER is None: GLOBAL_KEEP_MODEL_LOADER = KEEPModelLoader() return GLOBAL_KEEP_MODEL_LOADER -class KEEP_ModelLoaderNode: - _MODEL_TYPES = list(KEEP_MODEL_CONFIGS.keys()) +class KEEP_ModelLoaderNode(io.ComfyNode): + """Load the KEEP checkpoint, the face helper and the optional upscalers.""" @classmethod - def INPUT_TYPES(s): - detection_models = ['retinaface_resnet50', 'retinaface_mobile0.25', 'YOLOv5l', 'YOLOv5n'] - return { - "required": { - "model": (s._MODEL_TYPES, {"default": s._MODEL_TYPES[0] if s._MODEL_TYPES else 'KEEP'}), - "detection_model": (detection_models, {"default": 'retinaface_resnet50'}), - }, - "optional": { - "bg_upscale_model": ("UPSCALE_MODEL",), - "face_upscale_model": ("UPSCALE_MODEL",), - } - } - RETURN_TYPES = ("KEEP_MODEL_PACK",) - RETURN_NAMES = ("keep_model_pack",) - FUNCTION = "load_model_pack" - CATEGORY = "ComfyUI-KEEP" + def define_schema(cls) -> io.Schema: + model_types = list(KEEP_MODEL_CONFIGS.keys()) + return io.Schema( + node_id="KEEP_ModelLoader", + display_name="Load KEEP Models", + category="ComfyUI-KEEP", + description=( + "Load the KEEP face restoration model together with the face " + "detection / parsing models. The result is a model pack for " + "the KEEP Single Image and KEEP Image Sequence nodes." + ), + inputs=[ + io.Combo.Input( + "model", + options=model_types, + default=model_types[0] if model_types else "KEEP", + tooltip="The KEEP checkpoint to load (KEEP or Asian).", + ), + io.Combo.Input( + "detection_model", + options=list(DETECTION_MODELS), + default="retinaface_resnet50", + tooltip="The face detector of the face restoration helper.", + ), + io.Custom("UPSCALE_MODEL").Input( + "bg_upscale_model", + optional=True, + tooltip=( + "An optional ComfyUI upscale model for the background " + "of the input image." + ), + ), + io.Custom("UPSCALE_MODEL").Input( + "face_upscale_model", + optional=True, + tooltip=( + "An optional ComfyUI upscale model for the restored " + "face patch." + ), + ), + ], + outputs=[ + io.Custom("KEEP_MODEL_PACK").Output( + id="keep_model_pack", + display_name="keep_model_pack", + tooltip="The loaded KEEP model pack.", + ), + ], + ) - def load_model_pack(self, model, detection_model, bg_upscale_model=None, face_upscale_model=None): + @classmethod + def execute( + cls, + model, + detection_model, + bg_upscale_model=None, + face_upscale_model=None, + ) -> io.NodeOutput: loader = get_keep_model_loader() model_pack = loader.load_keep_model_pack( model_type_str=model, detection_model_str=detection_model, bg_upscale_model=bg_upscale_model, - face_upscale_model=face_upscale_model + face_upscale_model=face_upscale_model, ) - return (model_pack,) + return io.NodeOutput(model_pack) diff --git a/tests/test_nodes_execute.py b/tests/test_nodes_execute.py index 36cf0d8..8c3357b 100644 --- a/tests/test_nodes_execute.py +++ b/tests/test_nodes_execute.py @@ -1,10 +1,10 @@ -"""P3.2 - the mock-integration execution of the node classes. +"""P3.2 / P4.1 - the mock-integration execution of the node classes. -The nodes are still V1 classes in this phase (P4 migrates them to the V3 -schema), so their ``execute`` path is the plain method. The tests drive it -with a **stub pack** and a stubbed processor and assert the contract of the -node: the image that goes in comes back out, and the pack is loaded and -offloaded exactly once. +The loader node is a V3 node (P4.1) - its ``execute`` is a classmethod that +returns an ``io.NodeOutput``. The tests drive it with a **stubbed loader** +(the model build is stubbed out - no weights are needed) and assert the +contract of the node: the pack comes back, the widgets reach the loader +unchanged and the second call is served from the loader's cache. """ from __future__ import annotations @@ -26,6 +26,11 @@ def loader_module(): return importlib.import_module("ComfyUI_KEEP.src.keep_model_loader") +@pytest.fixture(scope="module") +def loader_node_module(): + return importlib.import_module("ComfyUI_KEEP.nodes.keep_model_loader") + + @pytest.fixture def stub_pack(loader_module, monkeypatch): """A real ``KEEPModelPack`` (the node checks the type) with stubs.""" @@ -43,6 +48,128 @@ def stub_pack(loader_module, monkeypatch): return pack +@pytest.fixture +def stubbed_loader(loader_module, monkeypatch, tmp_path): + """A real ``KEEPModelLoader`` whose model build is stubbed out. + + The stub replaces the four heavy parts of the build (the architecture, the + checkpoint download, the checkpoint load and the face helper) - the cache + logic of the loader itself stays **real**. + """ + class _StubArch: + def __init__(self, **kwargs): + self.kwargs = kwargs + self.loaded = None + + def to(self, device): + return self + + def load_state_dict(self, state_dict, strict=True): + self.loaded = state_dict + return self + + def eval(self): + return self + + class _StubHelper: + def __init__(self, **kwargs): + self.kwargs = kwargs + + monkeypatch.setattr(loader_module, "ARCH_REGISTRY", {"KEEP": _StubArch}) + monkeypatch.setattr( + loader_module, + "load_file_from_url_comfy", + lambda **kwargs: str(tmp_path / "checkpoint.pth"), + ) + monkeypatch.setattr( + loader_module.torch, + "load", + lambda *args, **kwargs: {"params_ema": {"weight": torch.zeros(1)}}, + ) + monkeypatch.setattr(loader_module, "_face_restore_helper", lambda: _StubHelper) + monkeypatch.setattr( + loader_module.keep_paths, "resolve_model_root", lambda *args, **kwargs: str(tmp_path) + ) + return loader_module.KEEPModelLoader() + + +# --------------------------------------------------------------------------- +# P4.1 - the model loader node +# --------------------------------------------------------------------------- +@pytest.mark.mock_integration +def test_the_loader_node_returns_the_pack(loader_node_module, stubbed_loader, monkeypatch): + """The node returns the pack of the loader as a ``NodeOutput``.""" + from comfy_api.latest import io + + monkeypatch.setattr(loader_node_module, "GLOBAL_KEEP_MODEL_LOADER", stubbed_loader) + + output = loader_node_module.KEEP_ModelLoaderNode.execute( + model="KEEP", + detection_model="retinaface_resnet50", + ) + + assert isinstance(output, io.NodeOutput) + assert len(output.result) == 1 + pack = output.result[0] + assert pack.model_type_str == "KEEP" + assert pack.keep_net is not None + + +@pytest.mark.mock_integration +def test_the_loader_node_serves_the_second_call_from_the_cache( + loader_node_module, stubbed_loader, monkeypatch +): + """The second call with the same widgets reuses the cached models.""" + monkeypatch.setattr(loader_node_module, "GLOBAL_KEEP_MODEL_LOADER", stubbed_loader) + + first = loader_node_module.KEEP_ModelLoaderNode.execute( + model="KEEP", detection_model="retinaface_resnet50" + ).result[0] + second = loader_node_module.KEEP_ModelLoaderNode.execute( + model="KEEP", detection_model="retinaface_resnet50" + ).result[0] + + assert first is not second, "every call returns its own pack" + assert first.keep_net is second.keep_net, "the models are cached" + assert first.face_helper is second.face_helper, "the models are cached" + assert len(stubbed_loader.loaded_models) == 1, stubbed_loader.loaded_models + + +@pytest.mark.mock_integration +def test_the_loader_node_passes_the_widgets_to_the_loader( + loader_node_module, monkeypatch +): + """The widget values reach the loader unchanged.""" + received: list[dict] = [] + + class _RecordingLoader: + def load_keep_model_pack(self, **kwargs): + received.append(kwargs) + return "pack" + + monkeypatch.setattr(loader_node_module, "GLOBAL_KEEP_MODEL_LOADER", _RecordingLoader()) + + output = loader_node_module.KEEP_ModelLoaderNode.execute( + model="KEEP", + detection_model="retinaface_mobile0.25", + bg_upscale_model="bg", + face_upscale_model="face", + ) + + assert output.result == ("pack",) + assert received == [ + { + "model_type_str": "KEEP", + "detection_model_str": "retinaface_mobile0.25", + "bg_upscale_model": "bg", + "face_upscale_model": "face", + } + ] + + +# --------------------------------------------------------------------------- +# P3.2 - the single image node (V1 until P4.2) +# --------------------------------------------------------------------------- @pytest.mark.mock_integration def test_upscale_face_image_returns_the_processed_image(node_module, stub_pack, monkeypatch): """The node returns the processed image and uses the pack once.""" diff --git a/tests/test_nodes_layout.py b/tests/test_nodes_layout.py index 4d70e8c..9710dc2 100644 --- a/tests/test_nodes_layout.py +++ b/tests/test_nodes_layout.py @@ -1,15 +1,14 @@ -"""P3.1 - the node classes live in ``nodes/``, one file per node. +"""P4.1 - the V3 schema of the nodes keeps the compatibility contract. -P3.1 is a **pure move**: the V1 classes of the former root ``nodes.py`` are -moved verbatim into ``nodes/``. The tests below therefore assert the -*compatibility contract* of the pack against the **live** classes: +The node ids, the display names, the category, the socket types and every +widget (name, default, min/max/step and order) are compatibility guarantees +of the pack: existing user workflows must keep working unchanged. -* the node ids, the display names and the category, -* the socket types (IMAGE / KEEP_MODEL_PACK / UPSCALE_MODEL), -* every widget: its name, its default, its min/max/step and its order. - -A move that changes any of these would break existing user workflows, so the -snapshot is the real acceptance criterion of this step. +The contract below is the **recorded V1 ``INPUT_TYPES``** (P0.1). After the +V3 migration the widgets live in the schema, so the test converts the V3 +schema into the very same shape and compares it with the record. A V3 +schema that drops a widget, renames one, changes a default or reorders the +inputs fails here. """ from __future__ import annotations @@ -38,7 +37,11 @@ PACK_DISPLAY_NAMES = { "KEEP_ProcessImageSequence": "KEEP Image Sequence", } -#: The expected widget contract of every node (V1 ``INPUT_TYPES``). +#: The category of the pack - it must never change. +PACK_CATEGORY = "ComfyUI-KEEP" + +#: The expected widget contract of every node (the recorded V1 +#: ``INPUT_TYPES`` - the source of truth of this file). EXPECTED_INPUT_TYPES = { "KEEP_ModelLoader": { "required": { @@ -144,34 +147,154 @@ EXPECTED_INPUT_TYPES = { }, } -#: The V1 node metadata (return types, function name, category). -EXPECTED_NODE_METADATA = { +#: The expected output sockets of every node (the recorded V1 +#: ``RETURN_TYPES`` / ``RETURN_NAMES``). +EXPECTED_OUTPUTS = { "KEEP_ModelLoader": { - "RETURN_TYPES": ("KEEP_MODEL_PACK",), - "RETURN_NAMES": ("keep_model_pack",), - "FUNCTION": "load_model_pack", - "CATEGORY": "ComfyUI-KEEP", + "types": ("KEEP_MODEL_PACK",), + "names": ("keep_model_pack",), }, "KEEP_FaceUpscaleImage": { - "RETURN_TYPES": ("IMAGE",), - "FUNCTION": "upscale_face_image", - "CATEGORY": "ComfyUI-KEEP", + "types": ("IMAGE",), + # the V1 node had no RETURN_NAMES - the socket type is the name + "names": ("IMAGE",), }, "KEEP_ProcessImageSequence": { - "RETURN_TYPES": ("IMAGE",), - "RETURN_NAMES": ("processed_images",), - "FUNCTION": "process_sequence", - "CATEGORY": "ComfyUI-KEEP", + "types": ("IMAGE",), + "names": ("processed_images",), }, } +# --------------------------------------------------------------------------- +# helpers +# --------------------------------------------------------------------------- +#: The class of every node id (the names are stable - they are the class +#: names of the pack). +NODE_CLASS_NAMES = { + "KEEP_ModelLoader": "KEEP_ModelLoaderNode", + "KEEP_FaceUpscaleImage": "KEEP_FaceUpscaleImageNode", + "KEEP_ProcessImageSequence": "KEEP_ProcessImageSequenceNode", +} + +#: The nodes that are **already** migrated to the V3 schema. The list grows +#: with the migration steps (P4.1 the loader, P4.2 the single image node, +#: P4.3 the sequence node); the contract test below works with both +#: representations, so it stays meaningful in every step. +V3_NODE_IDS = ("KEEP_ModelLoader",) + + +def _pack_nodes(nodes_package) -> dict: + """``{node_id: node class}`` - read from the exports of ``nodes/``.""" + return { + node_id: getattr(nodes_package, class_name) + for node_id, class_name in NODE_CLASS_NAMES.items() + } + + +def _is_v3(node_class) -> bool: + from comfy_api.latest import io + + return isinstance(node_class, type) and issubclass(node_class, io.ComfyNode) + + +def _v1_contract(input_types: dict) -> dict: + """The recorded V1 ``INPUT_TYPES`` in the shape of the node contract.""" + contract: dict[str, list[dict]] = {"required": [], "optional": []} + for section, widgets in input_types.items(): + for name, spec in widgets.items(): + io_type = spec[0] + options = spec[1] if len(spec) > 1 else {} + entry = { + "id": name, + "type": io_type if isinstance(io_type, str) else "COMBO", + "tooltip": options.get("tooltip"), + } + if not isinstance(io_type, str): + entry["options"] = list(io_type) + for field in ("default", "min", "max", "step"): + if field in options: + entry[field] = options[field] + contract[section].append(entry) + return contract + + +def _schema_contract(schema) -> dict: + """The V3 schema in the shape of the node contract.""" + contract: dict[str, list[dict]] = {"required": [], "optional": []} + for item in schema.inputs: + section = "optional" if item.optional else "required" + entry = { + "id": item.id, + "type": item.get_io_type(), + "tooltip": item.tooltip, + } + for field in ("default", "min", "max", "step", "options"): + value = getattr(item, field, None) + if value is not None: + entry[field] = value + contract[section].append(entry) + return contract + + +def _node_contract(node_class) -> dict: + """The contract of a node - from its V3 schema or from its V1 class. + + The V3 migration moves the widgets from ``INPUT_TYPES`` into the + schema; the contract of the pack is the same either way, so the test + compares the *contract*, not the representation. + """ + if _is_v3(node_class): + schema = node_class.GET_SCHEMA() + return { + "v3": True, + "node_id": schema.node_id, + "display_name": schema.display_name, + "category": schema.category, + "description": schema.description, + "inputs": _schema_contract(schema), + "output_types": tuple(output.get_io_type() for output in schema.outputs), + "output_names": tuple( + output.display_name if output.display_name else output.get_io_type() + for output in schema.outputs + ), + "tooltips": {item.id: item.tooltip for item in schema.inputs}, + } + input_types = node_class.INPUT_TYPES() + return { + "v3": False, + "node_id": None, + "display_name": None, + "category": node_class.CATEGORY, + "description": getattr(node_class, "DESCRIPTION", ""), + "inputs": _v1_contract(input_types), + "output_types": tuple(node_class.RETURN_TYPES), + "output_names": tuple( + getattr(node_class, "RETURN_NAMES", node_class.RETURN_TYPES) + ), + "tooltips": { + entry["id"]: entry["tooltip"] + for entries in _v1_contract(input_types).values() + for entry in entries + }, + } + + @pytest.fixture(scope="module") def nodes_package(): """The ``nodes`` package of the node pack (imported under the alias).""" return importlib.import_module("ComfyUI_KEEP.nodes") +@pytest.fixture(scope="module") +def pack_nodes(nodes_package): + """``{node_id: node class}`` of the three pack nodes.""" + return _pack_nodes(nodes_package) + + +# --------------------------------------------------------------------------- +# the layout +# --------------------------------------------------------------------------- @pytest.mark.unit def test_every_node_has_its_own_file(): """One file per node, and the old root ``nodes.py`` is gone.""" @@ -188,54 +311,103 @@ def test_every_node_has_its_own_file(): @pytest.mark.unit -def test_node_class_mappings_has_exactly_the_three_ids(nodes_package): - """``NODE_CLASS_MAPPINGS`` holds exactly the three pack ids.""" - mappings = nodes_package.NODE_CLASS_MAPPINGS - assert sorted(mappings) == sorted(PACK_NODE_IDS), sorted(mappings) - for node_id, node_class in mappings.items(): - assert isinstance(node_class, type), node_id - assert node_class.__name__.endswith("Node"), node_class.__name__ +def test_the_pack_registers_exactly_the_three_node_ids(pack_nodes): + """The three node ids of the pack - and no other node.""" + assert sorted(pack_nodes) == sorted(PACK_NODE_IDS), sorted(pack_nodes) @pytest.mark.unit -def test_display_name_mappings_are_unchanged(nodes_package): - """The display names of the three nodes are unchanged.""" - assert nodes_package.NODE_DISPLAY_NAME_MAPPINGS == PACK_DISPLAY_NAMES +def test_every_node_defines_a_v3_schema_and_execute(pack_nodes): + """Every migrated node is a V3 node: ``define_schema`` + ``execute``.""" + for node_id in V3_NODE_IDS: + node_class = pack_nodes[node_id] + assert _is_v3(node_class), f"{node_id} is not a V3 node yet" + schema = node_class.GET_SCHEMA() + assert schema.node_id == node_id, node_id + assert callable(node_class.execute), node_id + + +# --------------------------------------------------------------------------- +# the schema contract +# --------------------------------------------------------------------------- +@pytest.mark.unit +def test_the_display_names_and_the_category_are_unchanged(pack_nodes, nodes_package): + """The display names of the three nodes and the category are unchanged.""" + v1_display_names = getattr(nodes_package, "NODE_DISPLAY_NAME_MAPPINGS", {}) + for node_id, node_class in pack_nodes.items(): + contract = _node_contract(node_class) + display_name = contract["display_name"] or v1_display_names.get(node_id) + assert display_name == PACK_DISPLAY_NAMES[node_id], node_id + assert contract["category"] == PACK_CATEGORY, node_id + if contract["v3"]: + assert contract["description"], f"{node_id} needs a schema description" @pytest.mark.unit -def test_the_widget_contract_is_unchanged(nodes_package): - """Every widget keeps its name, default, min/max/step and order.""" +def test_the_widget_contract_is_unchanged(pack_nodes): + """Every widget keeps its name, type, default, min/max/step and order.""" for node_id, expected in EXPECTED_INPUT_TYPES.items(): - node_class = nodes_package.NODE_CLASS_MAPPINGS[node_id] - actual = node_class.INPUT_TYPES() - assert list(actual) == list(expected), ( - f"{node_id}: the input sections changed ({list(actual)} != {list(expected)})" + contract = _node_contract(pack_nodes[node_id]) + actual = contract["inputs"] + expected_contract = _v1_contract(expected) + + assert list(actual) == list(expected_contract), ( + f"{node_id}: the input sections changed ({list(actual)} != " + f"{list(expected_contract)})" ) - for section, widgets in expected.items(): - assert list(actual[section]) == list(widgets), ( - f"{node_id}/{section}: the widget order changed " - f"({list(actual[section])} != {list(widgets)})" - ) - for name, spec in widgets.items(): - assert actual[section][name] == spec, ( - f"{node_id}/{section}/{name}: {actual[section][name]} != {spec}" - ) + for section in expected_contract: + actual_entries = actual[section] + expected_entries = expected_contract[section] + assert [entry["id"] for entry in actual_entries] == [ + entry["id"] for entry in expected_entries + ], f"{node_id}/{section}: the widget order changed" + for actual_entry, expected_entry in zip(actual_entries, expected_entries): + for key, value in expected_entry.items(): + if key == "tooltip" and value is None: + # the V1 node had no tooltip here - the V3 schema + # requires one (the plan's P4.1: every input gets a + # tooltip), but the text is new + if contract["v3"]: + assert actual_entry["tooltip"], ( + f"{node_id}/{actual_entry['id']}: " + "the input needs a tooltip" + ) + continue + assert actual_entry.get(key) == value, ( + f"{node_id}/{actual_entry['id']}/{key}: " + f"{actual_entry.get(key)!r} != {value!r}" + ) @pytest.mark.unit -def test_the_node_metadata_is_unchanged(nodes_package): - """The socket types, the function names and the category are unchanged.""" - for node_id, expected in EXPECTED_NODE_METADATA.items(): - node_class = nodes_package.NODE_CLASS_MAPPINGS[node_id] - for attribute, value in expected.items(): - assert getattr(node_class, attribute) == value, ( - f"{node_id}.{attribute}: {getattr(node_class, attribute)!r} != {value!r}" +def test_every_input_has_a_tooltip(pack_nodes): + """Every input of every V3 node carries a tooltip.""" + for node_id, node_class in pack_nodes.items(): + if not _is_v3(node_class): + continue + schema = node_class.GET_SCHEMA() + assert schema.inputs, node_id + for item in schema.inputs: + assert isinstance(item.tooltip, str) and item.tooltip.strip(), ( + f"{node_id}/{item.id}: the input has no tooltip" ) @pytest.mark.unit -def test_the_example_workflow_ids_exist_in_the_mappings(nodes_package): +def test_the_output_sockets_are_unchanged(pack_nodes): + """The output socket types and their names are unchanged.""" + for node_id, expected in EXPECTED_OUTPUTS.items(): + contract = _node_contract(pack_nodes[node_id]) + assert contract["output_types"] == expected["types"], ( + f"{node_id}: {contract['output_types']} != {expected['types']}" + ) + assert contract["output_names"] == expected["names"], ( + f"{node_id}: {contract['output_names']} != {expected['names']}" + ) + + +@pytest.mark.unit +def test_the_example_workflow_ids_exist_in_the_pack(pack_nodes): """The ids of the shipped example are registered by the node.""" workflow = json.loads( (EXAMPLES_DIR / "keep-workflow.json").read_text(encoding="utf-8") @@ -246,9 +418,7 @@ def test_the_example_workflow_ids_exist_in_the_mappings(nodes_package): if isinstance(node.get("type"), str) and node["type"].startswith("KEEP_") } assert used, "the example must use the KEEP nodes" - assert used <= set(nodes_package.NODE_CLASS_MAPPINGS), sorted( - used - set(nodes_package.NODE_CLASS_MAPPINGS) - ) + assert used <= set(pack_nodes), sorted(used - set(pack_nodes)) @pytest.mark.unit