test: cover workflow components
This commit is contained in:
Dariusz L
@@ -0,0 +1,216 @@
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
from core.workflow import references, subgraphs
|
||||
|
||||
|
||||
def _loader_node(node_id=10, value="model.safetensors", link=None):
|
||||
return {
|
||||
"id": node_id,
|
||||
"type": "CheckpointLoaderSimple",
|
||||
"title": "Loader",
|
||||
"widgets_values": [value],
|
||||
"inputs": [
|
||||
{
|
||||
"name": "ckpt_name",
|
||||
"widget": {"name": "ckpt_name"},
|
||||
"link": link,
|
||||
}
|
||||
],
|
||||
"outputs": [],
|
||||
}
|
||||
|
||||
|
||||
def test_subgraph_widget_indexes_support_named_inputs_and_links():
|
||||
node = {
|
||||
**_loader_node(),
|
||||
"inputs": [
|
||||
{"name": "model", "widget": None, "link": 1},
|
||||
{"name": "ckpt_name", "widget": {"name": "ckpt_name"}},
|
||||
],
|
||||
"widgets_values": ["model.safetensors"],
|
||||
}
|
||||
|
||||
assert subgraphs._get_widget_index_for_input_index(node, 0) is None
|
||||
assert subgraphs._get_widget_index_for_input_index(node, 1) == 0
|
||||
assert subgraphs._get_widget_index_by_name(node, "ckpt-name") == 0
|
||||
assert subgraphs._get_node_by_id([node], "10") is node
|
||||
|
||||
subgraph = {
|
||||
"inputs": [{"name": "ckpt_name", "linkIds": [7, "8"]}],
|
||||
}
|
||||
assert subgraphs._get_subgraph_input_link_ids(subgraph, "ckpt-name") == {
|
||||
"7",
|
||||
"8",
|
||||
}
|
||||
assert subgraphs._get_subgraph_input_link_ids(subgraph, "missing") == set()
|
||||
|
||||
|
||||
def test_find_promoted_targets_prefers_explicit_node_then_linked_input():
|
||||
explicit_node = _loader_node(node_id=10)
|
||||
linked_node = _loader_node(node_id=11, link=7)
|
||||
subgraph = {
|
||||
"inputs": [{"name": "ckpt_name", "linkIds": [7]}],
|
||||
"nodes": [explicit_node, linked_node],
|
||||
}
|
||||
|
||||
with patch.object(
|
||||
subgraphs.references,
|
||||
"get_effective_model_category_hint",
|
||||
return_value="checkpoints",
|
||||
):
|
||||
explicit = subgraphs._find_promoted_widget_targets(
|
||||
subgraph, "10", "ckpt_name"
|
||||
)
|
||||
linked = subgraphs._find_promoted_widget_targets(
|
||||
subgraph, "-1", "ckpt_name"
|
||||
)
|
||||
|
||||
assert [(item["node_id"], item["widget_index"]) for item in explicit] == [(10, 0)]
|
||||
assert [(item["node_id"], item["widget_index"]) for item in linked] == [(11, 0)]
|
||||
assert linked[0]["category"] == "checkpoints"
|
||||
|
||||
|
||||
def test_build_promoted_contexts_records_instance_and_inner_locators():
|
||||
workflow = {
|
||||
"nodes": [
|
||||
{
|
||||
"id": 1,
|
||||
"type": "subgraph-1",
|
||||
"widgets_values": ["promoted.safetensors"],
|
||||
"properties": {"proxyWidgets": [["10", "ckpt_name"]]},
|
||||
}
|
||||
]
|
||||
}
|
||||
subgraph = {
|
||||
"id": "subgraph-1",
|
||||
"name": "Promoted loader",
|
||||
"nodes": [_loader_node()],
|
||||
}
|
||||
|
||||
with patch.object(
|
||||
subgraphs.references,
|
||||
"get_effective_model_category_hint",
|
||||
return_value="checkpoints",
|
||||
):
|
||||
contexts = subgraphs._build_promoted_widget_contexts(workflow, [subgraph])
|
||||
|
||||
instance = contexts["instance_widgets"][("1", 0)]
|
||||
assert instance["subgraph_name"] == "Promoted loader"
|
||||
assert instance["promoted_value"] == "promoted.safetensors"
|
||||
assert contexts["inner_widgets"][("subgraph-1", "10", 0)][0]["node_id"] == 1
|
||||
assert contexts["inner_widget_names"][("subgraph-1", "10", "ckpt_name")]
|
||||
|
||||
|
||||
def test_apply_instance_promoted_context_resolves_existing_model(tmp_path):
|
||||
model_path = tmp_path / "promoted.safetensors"
|
||||
model_path.write_bytes(b"model")
|
||||
available_models = [
|
||||
{
|
||||
"category": "checkpoints",
|
||||
"relative_path": "promoted.safetensors",
|
||||
"path": str(model_path),
|
||||
}
|
||||
]
|
||||
reference = {"original_path": "promoted.safetensors", "exists": False}
|
||||
context = {
|
||||
"proxy_widget_name": "ckpt_name",
|
||||
"node_id": 10,
|
||||
"node_type": "CheckpointLoaderSimple",
|
||||
"node_title": "Loader",
|
||||
"widget_index": 0,
|
||||
"widget_name": "ckpt_name",
|
||||
"category": "checkpoints",
|
||||
}
|
||||
|
||||
subgraphs._apply_instance_promoted_widget_context(
|
||||
reference,
|
||||
context,
|
||||
available_models,
|
||||
)
|
||||
|
||||
assert reference["promoted_inner_node_id"] == 10
|
||||
assert reference["category"] == "checkpoints"
|
||||
assert reference["full_path"] == str(model_path)
|
||||
assert reference["exists"] is True
|
||||
|
||||
|
||||
def test_apply_promoted_locator_supports_name_fallback_and_value_selection():
|
||||
contexts = {
|
||||
"inner_widgets": {},
|
||||
"inner_widget_names": {
|
||||
("subgraph-1", "10", "ckpt_name"): [
|
||||
{"node_id": 1, "node_type": "Subgraph", "subgraph_name": "One", "promoted_value": "first"},
|
||||
{"node_id": 2, "node_type": "Subgraph", "subgraph_name": "Two", "promoted_value": "second"},
|
||||
]
|
||||
},
|
||||
}
|
||||
reference = {
|
||||
"subgraph_id": "subgraph-1",
|
||||
"node_id": 10,
|
||||
"widget_index": 99,
|
||||
"widget_name": "ckpt_name",
|
||||
"original_path": "second",
|
||||
}
|
||||
|
||||
subgraphs._apply_promoted_widget_locator(reference, contexts)
|
||||
|
||||
assert reference["locate_node_id"] == 2
|
||||
assert reference["locate_subgraph_name"] == "Two"
|
||||
assert reference["locate_is_top_level"] is True
|
||||
assert reference["locate_via_promoted_widget"] is True
|
||||
|
||||
|
||||
def test_reference_matching_handles_placeholders_paths_and_scanner_records(tmp_path):
|
||||
model_path = tmp_path / "folder" / "model.safetensors"
|
||||
model_path.parent.mkdir()
|
||||
model_path.write_bytes(b"model")
|
||||
available_models = [
|
||||
{
|
||||
"category": "checkpoints",
|
||||
"relative_path": "folder/model.safetensors",
|
||||
"path": str(model_path),
|
||||
}
|
||||
]
|
||||
|
||||
assert references.is_model_filename("model.safetensors") is True
|
||||
assert references.is_model_filename("urn:air:sd:checkpoint:civitai:1@2") is True
|
||||
assert references.should_scan_as_model_reference("None", True) is False
|
||||
assert references.should_scan_as_model_reference("plain-name", True) is True
|
||||
assert references.should_scan_as_model_reference("plain-name", False) is False
|
||||
assert references.is_static_or_hybrid_widget_choice(
|
||||
"Folder\\Model.safetensors",
|
||||
{"source": "static", "choices": ["folder/model.safetensors"]},
|
||||
)
|
||||
assert references.static_or_hybrid_choice_looks_like_model(
|
||||
"folder/model.safetensors",
|
||||
{"source": "hybrid", "choices": []},
|
||||
)
|
||||
|
||||
resolved = references.try_resolve_model_path(
|
||||
"folder/model.safetensors",
|
||||
["checkpoints", "configs"],
|
||||
available_models=available_models,
|
||||
)
|
||||
assert resolved == ("checkpoints", str(model_path))
|
||||
assert references.try_resolve_model_path(
|
||||
"folder/MODEL.safetensors",
|
||||
["checkpoints"],
|
||||
available_models=available_models,
|
||||
) is None
|
||||
|
||||
|
||||
def test_reference_resolution_uses_folder_paths_when_scanner_is_not_available(tmp_path):
|
||||
model_path = tmp_path / "model.safetensors"
|
||||
model_path.write_bytes(b"model")
|
||||
fake_folder_paths = SimpleNamespace(
|
||||
folder_names_and_paths={"checkpoints": ([], set()), "configs": ([], set())},
|
||||
get_filename_list=lambda category: ["model.safetensors"],
|
||||
get_full_path=lambda category, filename: str(model_path),
|
||||
)
|
||||
|
||||
with patch.object(references, "folder_paths", fake_folder_paths):
|
||||
assert references.try_resolve_model_path("model.safetensors") == (
|
||||
"checkpoints",
|
||||
str(model_path),
|
||||
)
|
||||
Reference in New Issue
Block a user