nodes/wildminder-ComfyUI-KEEP
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

plan(P4.1): migrate the model loader node to the V3 schema

Browse Source
This commit is contained in:
WildAi
2026-10-23 20:55:15 +03:00
parent fcbfcce3a6
commit c49f3832af
4 changed files with 501 additions and 150 deletions
Download Patch File Download Diff File Expand all files Collapse all files
conftest.py
+61 -59
View File
@@ -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:
nodes/keep_model_loader.py
+77 -25
View File
@@ -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)
tests/test_nodes_execute.py
+133 -6
View File
@@ -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."""
tests/test_nodes_layout.py
+230 -60
View File
@@ -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