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), )