373 lines
15 KiB
Python
373 lines
15 KiB
Python
"""P6.1 - the shipped example workflows are valid against the V3 node set.
|
|
|
|
Every workflow json in ``example_workflows/`` must
|
|
|
|
* parse (it is the ComfyUI UI format, version 0.4),
|
|
* use only node types that ComfyUI really has: a **pack node id** (the ids are
|
|
regex-parsed from ``nodes/*.py``, so a renamed or removed node is caught),
|
|
a **whitelisted core node** or a **whitelisted VideoHelperSuite type**
|
|
(``VHS_LoadVideo`` / ``VHS_VideoCombine`` - the shipped example is a video
|
|
workflow and therefore needs the VideoHelperSuite pack; the dependency is
|
|
documented in the README),
|
|
* have a preview image next to it (the file the registry and the README show),
|
|
* have consistent links (every link points at a node that exists, every
|
|
connected input refers to a link that exists), and
|
|
* match the **V3 schema** of the pack nodes they use: the same number of
|
|
widget values (in the same order), the same socket names and types, and
|
|
widget values that are legal for their widget (a combo value must be one of
|
|
its options, a number must be inside its range).
|
|
|
|
The last group is what makes "the shipped example still loads with the
|
|
migrated ids" a checked statement instead of a claim: a workflow whose widget
|
|
values no longer fit the V3 schema would be rejected by ComfyUI.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib
|
|
import json
|
|
import re
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parent.parent
|
|
EXAMPLES_DIR = REPO_ROOT / "example_workflows"
|
|
NODES_DIR = REPO_ROOT / "nodes"
|
|
|
|
#: The core ComfyUI node types the examples may use.
|
|
CORE_NODE_WHITELIST = {
|
|
"LoadImage",
|
|
"PreviewImage",
|
|
"SaveImage",
|
|
"UpscaleModelLoader",
|
|
"RepeatImageBatch",
|
|
"ImageBatch",
|
|
}
|
|
|
|
#: The VideoHelperSuite node types the shipped example uses. They are a
|
|
#: **dependency of the example**, not of the pack: the README documents that
|
|
#: the video example needs the VideoHelperSuite pack.
|
|
VHS_NODE_WHITELIST = {
|
|
"VHS_LoadVideo",
|
|
"VHS_VideoCombine",
|
|
}
|
|
|
|
#: The node ids the pack registers - the compatibility guarantee of P4.
|
|
EXPECTED_PACK_NODE_IDS = {
|
|
"KEEP_ModelLoader",
|
|
"KEEP_FaceUpscaleImage",
|
|
"KEEP_ProcessImageSequence",
|
|
}
|
|
|
|
#: The socket types that are **widgets** in the V3 schema (they are the
|
|
#: ``widgets_values`` of a saved workflow).
|
|
WIDGET_SOCKET_TYPES = {"INT", "FLOAT", "STRING", "BOOLEAN", "COMBO", "MULTICOMBO"}
|
|
|
|
#: The png magic bytes.
|
|
PNG_MAGIC = b"\x89PNG\r\n\x1a\n"
|
|
|
|
|
|
def _workflow_files() -> list[Path]:
|
|
return sorted(EXAMPLES_DIR.glob("*.json"))
|
|
|
|
|
|
def _load(path: Path) -> dict:
|
|
return json.loads(path.read_text(encoding="utf-8"))
|
|
|
|
|
|
def _pack_node_ids() -> set[str]:
|
|
"""The node ids the pack registers - parsed from the sources.
|
|
|
|
The regex is the one the sources use (``node_id="..."`` in the V3
|
|
schemas), so the whitelist of the test follows the code, not a copy of it.
|
|
"""
|
|
ids: set[str] = set()
|
|
for source in sorted(NODES_DIR.glob("*.py")):
|
|
ids.update(re.findall(r'node_id="([^"]+)"', source.read_text(encoding="utf-8")))
|
|
return ids
|
|
|
|
|
|
def _widget_inputs(schema) -> list:
|
|
return [i for i in schema.inputs if i.get_io_type() in WIDGET_SOCKET_TYPES]
|
|
|
|
|
|
def _socket_inputs(schema) -> list:
|
|
return [i for i in schema.inputs if i.get_io_type() not in WIDGET_SOCKET_TYPES]
|
|
|
|
|
|
def _output_name(output) -> str:
|
|
"""The name a saved workflow uses for an output of the schema."""
|
|
return getattr(output, "display_name", None) or getattr(output, "id", "")
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def pack_schemas() -> dict:
|
|
"""The V3 schema of every pack node, by node id."""
|
|
module = importlib.import_module("ComfyUI_KEEP.nodes")
|
|
schemas = {}
|
|
for class_name in module.__all__:
|
|
schema = getattr(module, class_name).define_schema()
|
|
schemas[schema.node_id] = schema
|
|
return schemas
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_the_examples_directory_holds_the_two_workflows():
|
|
"""The canonical examples directory holds both example workflows."""
|
|
names = {path.name for path in _workflow_files()}
|
|
assert "keep-workflow.json" in names, (
|
|
"the shipped example workflow must stay in example_workflows/"
|
|
)
|
|
assert "keep-sequence-workflow.json" in names, (
|
|
"P6.1 adds the sequence example next to the shipped one"
|
|
)
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_every_example_workflow_parses():
|
|
"""Every example is a parseable ComfyUI UI workflow."""
|
|
for path in _workflow_files():
|
|
data = _load(path)
|
|
assert isinstance(data, dict), f"{path.name} is not a workflow object"
|
|
assert isinstance(data.get("nodes"), list) and data["nodes"], (
|
|
f"{path.name} has no nodes"
|
|
)
|
|
assert isinstance(data.get("links"), list), f"{path.name} has no links list"
|
|
assert data.get("version") == 0.4, (
|
|
f"{path.name} is not in the UI format 0.4 (got {data.get('version')!r})"
|
|
)
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_the_pack_registers_the_three_migrated_node_ids():
|
|
"""The ids the examples whitelist are the ids the pack registers."""
|
|
assert _pack_node_ids() == EXPECTED_PACK_NODE_IDS, (
|
|
f"the pack registers {sorted(_pack_node_ids())}, "
|
|
f"the examples expect {sorted(EXPECTED_PACK_NODE_IDS)}"
|
|
)
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_every_node_type_is_a_pack_core_or_vhs_node():
|
|
"""No example uses a node type that does not exist."""
|
|
known = _pack_node_ids() | CORE_NODE_WHITELIST | VHS_NODE_WHITELIST
|
|
for path in _workflow_files():
|
|
data = _load(path)
|
|
types = {node.get("type") for node in data["nodes"]}
|
|
unknown = {t for t in types if t not in known}
|
|
assert not unknown, (
|
|
f"{path.name} uses node types that do not exist: {sorted(unknown)}"
|
|
)
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_every_workflow_has_its_preview_image():
|
|
"""Every example ships the preview image that belongs to it."""
|
|
for path in _workflow_files():
|
|
preview = path.with_suffix(".png")
|
|
assert preview.is_file(), f"the preview {preview.name} of {path.name} is missing"
|
|
data = preview.read_bytes()
|
|
assert data[:8] == PNG_MAGIC, f"{preview.name} is not a PNG file"
|
|
assert len(data) > 0, f"{preview.name} is empty"
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_the_links_of_every_workflow_are_consistent():
|
|
"""Every link of every example points at nodes and slots that exist."""
|
|
for path in _workflow_files():
|
|
data = _load(path)
|
|
nodes = {node["id"]: node for node in data["nodes"]}
|
|
assert len(nodes) == len(data["nodes"]), f"{path.name} has duplicate node ids"
|
|
|
|
link_ids = []
|
|
for link in data["links"]:
|
|
assert isinstance(link, list) and len(link) == 6, (
|
|
f"{path.name} has a malformed link: {link!r}"
|
|
)
|
|
link_id, source_id, source_slot, target_id, target_slot, kind = link
|
|
link_ids.append(link_id)
|
|
assert source_id in nodes, f"{path.name}: link {link_id} has no source node"
|
|
assert target_id in nodes, f"{path.name}: link {link_id} has no target node"
|
|
source = nodes[source_id]
|
|
target = nodes[target_id]
|
|
assert source_slot < len(source.get("outputs", [])), (
|
|
f"{path.name}: link {link_id} uses the output {source_slot} of "
|
|
f"{source.get('type')}, which has {len(source.get('outputs', []))}"
|
|
)
|
|
assert target_slot < len(target.get("inputs", [])), (
|
|
f"{path.name}: link {link_id} uses the input {target_slot} of "
|
|
f"{target.get('type')}, which has {len(target.get('inputs', []))}"
|
|
)
|
|
source_type = source["outputs"][source_slot].get("type")
|
|
target_type = target["inputs"][target_slot].get("type")
|
|
assert source_type == target_type == kind, (
|
|
f"{path.name}: link {link_id} connects {source_type} to {target_type} "
|
|
f"(the link says {kind})"
|
|
)
|
|
assert len(link_ids) == len(set(link_ids)), f"{path.name} has duplicate link ids"
|
|
|
|
# every connected input refers to a link that exists and points back
|
|
for node in data["nodes"]:
|
|
for index, socket in enumerate(node.get("inputs", [])):
|
|
link_id = socket.get("link")
|
|
if link_id is None:
|
|
continue
|
|
assert link_id in link_ids, (
|
|
f"{path.name}: the input {socket.get('name')} of node "
|
|
f"{node['id']} refers to the unknown link {link_id}"
|
|
)
|
|
link = next(l for l in data["links"] if l[0] == link_id)
|
|
assert link[3] == node["id"] and link[4] == index, (
|
|
f"{path.name}: the link {link_id} does not point at the input "
|
|
f"{index} of node {node['id']}"
|
|
)
|
|
|
|
# every output lists exactly the links that start at it
|
|
for node in data["nodes"]:
|
|
for slot, output in enumerate(node.get("outputs", [])):
|
|
listed = output.get("links") or []
|
|
expected = [
|
|
link[0]
|
|
for link in data["links"]
|
|
if link[1] == node["id"] and link[2] == slot
|
|
]
|
|
assert sorted(listed) == sorted(expected), (
|
|
f"{path.name}: the output {slot} of node {node['id']} lists the "
|
|
f"links {listed}, the graph has {expected}"
|
|
)
|
|
|
|
assert data.get("last_node_id", 0) >= max(nodes, default=0), (
|
|
f"{path.name}: last_node_id is smaller than the highest node id"
|
|
)
|
|
assert data.get("last_link_id", 0) >= max(link_ids, default=0), (
|
|
f"{path.name}: last_link_id is smaller than the highest link id"
|
|
)
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_the_pack_nodes_of_every_workflow_match_their_schema(pack_schemas):
|
|
"""The pack nodes of the examples fit the V3 schema of the pack.
|
|
|
|
This is the "the example still loads with the migrated ids" check: a
|
|
workflow with the wrong number of widget values, an unknown socket name or
|
|
an illegal widget value is rejected by ComfyUI.
|
|
"""
|
|
for path in _workflow_files():
|
|
data = _load(path)
|
|
for node in data["nodes"]:
|
|
schema = pack_schemas.get(node.get("type"))
|
|
if schema is None:
|
|
continue # a core or VHS node - not ours to check
|
|
|
|
widgets = _widget_inputs(schema)
|
|
values = node.get("widgets_values")
|
|
assert isinstance(values, list), (
|
|
f"{path.name}: the widget values of {node['type']} are not a list"
|
|
)
|
|
assert len(values) == len(widgets), (
|
|
f"{path.name}: {node['type']} has {len(values)} widget values, "
|
|
f"the schema has {len(widgets)} widgets"
|
|
)
|
|
|
|
for widget, value in zip(widgets, values):
|
|
kind = widget.get_io_type()
|
|
if kind == "COMBO":
|
|
assert value in (widget.options or []), (
|
|
f"{path.name}: {value!r} is not an option of the widget "
|
|
f"{widget.id} ({widget.options})"
|
|
)
|
|
elif kind == "BOOLEAN":
|
|
assert isinstance(value, bool), (
|
|
f"{path.name}: the widget {widget.id} of {node['type']} is "
|
|
f"not a boolean ({value!r})"
|
|
)
|
|
elif kind in {"INT", "FLOAT"}:
|
|
assert isinstance(value, (int, float)) and not isinstance(value, bool), (
|
|
f"{path.name}: the widget {widget.id} of {node['type']} is "
|
|
f"not a number ({value!r})"
|
|
)
|
|
if widget.min is not None:
|
|
assert value >= widget.min, (
|
|
f"{path.name}: the widget {widget.id} of {node['type']} is "
|
|
f"below its minimum ({value} < {widget.min})"
|
|
)
|
|
if widget.max is not None:
|
|
assert value <= widget.max, (
|
|
f"{path.name}: the widget {widget.id} of {node['type']} is "
|
|
f"above its maximum ({value} > {widget.max})"
|
|
)
|
|
elif kind == "STRING":
|
|
assert isinstance(value, str), (
|
|
f"{path.name}: the widget {widget.id} of {node['type']} is "
|
|
f"not a string ({value!r})"
|
|
)
|
|
|
|
sockets = {socket.id: socket for socket in _socket_inputs(schema)}
|
|
listed = {socket.get("name"): socket for socket in node.get("inputs", [])}
|
|
for name, socket in listed.items():
|
|
assert name in sockets, (
|
|
f"{path.name}: {node['type']} has the input {name!r}, which the "
|
|
f"schema does not know"
|
|
)
|
|
assert socket.get("type") == sockets[name].get_io_type(), (
|
|
f"{path.name}: the input {name!r} of {node['type']} is "
|
|
f"{socket.get('type')}, the schema says {sockets[name].get_io_type()}"
|
|
)
|
|
for socket in _socket_inputs(schema):
|
|
if not socket.optional:
|
|
assert socket.id in listed, (
|
|
f"{path.name}: {node['type']} misses the required input "
|
|
f"{socket.id!r}"
|
|
)
|
|
assert listed[socket.id].get("link") is not None, (
|
|
f"{path.name}: the required input {socket.id!r} of "
|
|
f"{node['type']} is not connected"
|
|
)
|
|
|
|
expected_outputs = {
|
|
_output_name(output): output.get_io_type() for output in schema.outputs
|
|
}
|
|
actual_outputs = {
|
|
output.get("name"): output.get("type") for output in node.get("outputs", [])
|
|
}
|
|
assert actual_outputs == expected_outputs, (
|
|
f"{path.name}: the outputs of {node['type']} are {actual_outputs}, "
|
|
f"the schema says {expected_outputs}"
|
|
)
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_the_shipped_example_uses_the_v3_ids_and_the_vhs_nodes():
|
|
"""The shipped example uses the migrated ids and the VHS dependency."""
|
|
data = _load(EXAMPLES_DIR / "keep-workflow.json")
|
|
types = {node.get("type") for node in data["nodes"]}
|
|
assert EXPECTED_PACK_NODE_IDS <= types, (
|
|
f"the shipped example must use the three pack nodes, it uses {sorted(types)}"
|
|
)
|
|
assert VHS_NODE_WHITELIST <= types, (
|
|
"the shipped example is a video workflow - it needs the VideoHelperSuite "
|
|
f"nodes {sorted(VHS_NODE_WHITELIST)}"
|
|
)
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_the_sequence_example_uses_the_sequence_node_and_png_previews():
|
|
"""The sequence example runs the sequence node and writes PNG previews."""
|
|
data = _load(EXAMPLES_DIR / "keep-sequence-workflow.json")
|
|
types = {node.get("type") for node in data["nodes"]}
|
|
assert "KEEP_ProcessImageSequence" in types, (
|
|
"the sequence example must use the KEEP_ProcessImageSequence node"
|
|
)
|
|
assert "KEEP_ModelLoader" in types, (
|
|
"the sequence example must load the model with the KEEP_ModelLoader node"
|
|
)
|
|
# SaveImage writes PNG files (the core node has no other format), the
|
|
# PreviewImage shows the frames in the UI.
|
|
assert "SaveImage" in types, (
|
|
"the sequence example must save the restored frames as PNG previews"
|
|
)
|
|
assert "PreviewImage" in types, (
|
|
"the sequence example must preview the restored frames"
|
|
)
|