422 lines
18 KiB
Python
422 lines
18 KiB
Python
"""CPU-only graph/compiler tests; no ComfyUI package bootstrap is needed."""
|
|
|
|
import copy
|
|
import importlib
|
|
import json
|
|
import subprocess
|
|
import sys
|
|
import types
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
PACKAGE = "serverless_media_prompt_testpkg"
|
|
# Real paths let sibling imports use production code, without running the
|
|
# extension's ComfyUI-dependent __init__.py or API registration side effects.
|
|
for suffix in ("", ".api", ".api.orchestration", ".serverless_bridge"):
|
|
package = types.ModuleType(PACKAGE + suffix)
|
|
package.__path__ = [str(ROOT.joinpath(*suffix.strip(".").split(".")))] if suffix else [str(ROOT)]
|
|
sys.modules[package.__name__] = package
|
|
|
|
compiler = importlib.import_module(f"{PACKAGE}.api.orchestration.serverless_media_prompt")
|
|
legacy = importlib.import_module(f"{PACKAGE}.api.orchestration.serverless_prompt")
|
|
contract = importlib.import_module(f"{PACKAGE}.serverless_bridge.contract")
|
|
compile_media_prompt = compiler.compile_media_prompt
|
|
SOCKETS = {"images": "IMAGE", "audio": "AUDIO", "masks": "MASK", "texts": "STRING", "video": "VIDEO"}
|
|
|
|
|
|
def _node(class_type, **inputs):
|
|
return {"class_type": class_type, "inputs": inputs}
|
|
|
|
|
|
def _collector(**sockets):
|
|
return _node(
|
|
"DistributedMediaCollector", multi_job_id="job/42", worker_id="worker_0", **sockets
|
|
)
|
|
|
|
|
|
def _prompt():
|
|
return {
|
|
"source:1": _node("LoadImage", image="example.png"),
|
|
"collector:1": _collector(images=["source:1", 0]),
|
|
}
|
|
|
|
|
|
def _run(prompt):
|
|
return compile_media_prompt(prompt, "request:1")
|
|
|
|
|
|
@pytest.mark.parametrize("socket,kind", SOCKETS.items())
|
|
def test_one_export_for_each_typed_socket_with_exact_canonical_binding(socket, kind):
|
|
prompt = {"source": _node("Source"), "collector": _collector(**{socket: ["source", 0]})}
|
|
graph, bindings = _run(prompt)
|
|
assert len(bindings) == 1
|
|
binding = bindings[0]
|
|
assert set(binding) == contract.BINDING_FIELDS
|
|
assert binding == {
|
|
"binding_id": binding["binding_id"],
|
|
"collector_id": "collector",
|
|
"job_id": "job/42",
|
|
"participant_id": "worker_0",
|
|
"node_id": binding["node_id"],
|
|
"socket": socket,
|
|
"kind": kind,
|
|
}
|
|
assert all(isinstance(value, str) for value in binding.values())
|
|
assert contract.validate_bindings(bindings) == {binding["binding_id"]: binding}
|
|
assert graph[binding["node_id"]] == {
|
|
"class_type": "MediaExportV1",
|
|
"inputs": {
|
|
"request_id": "request:1",
|
|
"binding_json": json.dumps(binding),
|
|
socket: ["source", 0],
|
|
},
|
|
}
|
|
assert set(graph) == {"source", binding["node_id"]}
|
|
|
|
|
|
def test_mixed_collectors_all_five_kinds_have_independent_dispatch_targets():
|
|
prompt = {
|
|
"shared": _node("MultiMediaSource"),
|
|
"collector:z": _collector(audio=["shared", 1], video=["shared", 4]),
|
|
"collector:a": _collector(images=["shared", 0], masks=["shared", 2], texts=["shared", 3]),
|
|
}
|
|
prompt["collector:z"]["inputs"].update(multi_job_id="job:z", worker_id="participant:z")
|
|
graph, bindings = _run(prompt)
|
|
assert [(b["collector_id"], b["socket"]) for b in bindings] == [
|
|
("collector:a", "images"), ("collector:a", "masks"), ("collector:a", "texts"),
|
|
("collector:z", "audio"), ("collector:z", "video"),
|
|
]
|
|
assert len({b["binding_id"] for b in bindings}) == 5
|
|
targets = {b["node_id"] for b in bindings}
|
|
assert len(targets) == 5
|
|
assert set(graph) == {"shared"} | targets
|
|
for binding in bindings:
|
|
export = graph[binding["node_id"]]
|
|
inputs = prompt[binding["collector_id"]]["inputs"]
|
|
assert export["inputs"][binding["socket"]] == inputs[binding["socket"]]
|
|
assert binding["job_id"] == inputs["multi_job_id"]
|
|
assert binding["participant_id"] == inputs["worker_id"]
|
|
|
|
|
|
def test_prunes_descendants_and_unrelated_outputs_preserving_required_output_utility():
|
|
prompt = _prompt()
|
|
prompt.update({
|
|
"required-output": _node("RequiredOutputUtility", images=["source:1", 0]),
|
|
"consumer": _node("ImageTransform", image=["required-output", 0]),
|
|
"descendant": _node("ImageTransform", image=["collector:1", 0]),
|
|
"after": _node("SaveImage", images=["descendant", 0]),
|
|
"unused-output": _node("SaveImage", images=["source:1", 0]),
|
|
"unrelated": _node("LoadAudio"),
|
|
})
|
|
prompt["required-output"]["_meta"] = {"title": "Required OUTPUT_NODE utility"}
|
|
prompt["collector:1"]["inputs"].update(images=["consumer", 0], load_balance=True, pass_through=False)
|
|
graph, bindings = _run(prompt)
|
|
assert set(graph) == {"source:1", "required-output", "consumer", bindings[0]["node_id"]}
|
|
assert graph["required-output"] == prompt["required-output"]
|
|
assert not {"load_balance", "pass_through"} & graph[bindings[0]["node_id"]]["inputs"].keys()
|
|
|
|
|
|
def test_stock_primitive_string_preserves_empty_and_link_looking_text():
|
|
for text in ("", '["source:1", 0]', "hello\nworld"):
|
|
prompt = {"text": _node("PrimitiveString", value=text), "collector": _collector(texts=["text", 0])}
|
|
graph, bindings = _run(prompt)
|
|
assert graph["text"] == prompt["text"]
|
|
assert graph[bindings[0]["node_id"]]["inputs"]["texts"] == ["text", 0]
|
|
|
|
|
|
@pytest.mark.parametrize("class_type,inputs", [
|
|
("DistributedSeed", {}),
|
|
("DistributedSeed", {"seed": 10, "is_worker": True, "worker_id": "worker_1"}),
|
|
("DistributedSeed", {"seed": 10, "is_worker": False, "worker_id": "worker_1"}),
|
|
("DistributedSeed", {"seed": 10, "is_worker": True, "worker_id": "bad"}),
|
|
("DistributedValue", {"default_value": "3.9", "worker_values": '{"_type":"INT","1":"8.7"}', "is_worker": True, "worker_id": "worker_0"}),
|
|
("DistributedValue", {"default_value": "3.9", "worker_values": '{"_type":"INT","2":""}', "is_worker": True, "worker_id": "worker_1"}),
|
|
("DistributedValue", {"default_value": "fallback", "worker_values": "malformed"}),
|
|
("DistributedValue", {"default_value": 1.25, "worker_values": '{"_type":"FLOAT","1":"bad"}', "is_worker": True, "worker_id": "0"}),
|
|
("DistributedModelName", {"text": "model.safetensors"}),
|
|
("DistributedModelName", {"text": {"name": "checkpoint"}}),
|
|
("DistributedModelName", {"text": ["a", "b", "c"]}),
|
|
])
|
|
def test_upstream_scalar_values_match_legacy_compiler(class_type, inputs):
|
|
prompt = _prompt()
|
|
prompt["scalar"] = _node(class_type, **inputs)
|
|
prompt["source:1"]["inputs"]["parameter"] = ["scalar", 0]
|
|
old_prompt = copy.deepcopy(prompt)
|
|
old_prompt["collector:1"]["class_type"] = "DistributedCollector"
|
|
expected, _ = legacy.compile_serverless_prompt(old_prompt)
|
|
graph, _ = _run(prompt)
|
|
assert graph["source:1"] == expected["source:1"]
|
|
assert "scalar" not in graph
|
|
|
|
|
|
def test_imports_same_scalar_resolver_and_resolves_primitive_chains():
|
|
assert compiler._make_scalar_resolver is legacy._make_scalar_resolver
|
|
prompt = _prompt()
|
|
prompt.update({
|
|
"int": _node("PrimitiveInt", value=25),
|
|
"int-alias": _node("PrimitiveInt", value=["int", 0]),
|
|
"float": _node("PrimitiveFloat", value=1.25),
|
|
"bool": _node("PrimitiveBoolean", value=True),
|
|
"string": _node("PrimitiveString", value="model.safetensors"),
|
|
"seed": _node("DistributedSeed", seed=["int-alias", 0], is_worker=True, worker_id="worker_0"),
|
|
"value": _node("DistributedValue", default_value=["float", 0], worker_values='{"_type":"FLOAT","1":"2.75"}', is_worker=True, worker_id="worker_0"),
|
|
"enabled": _node("DistributedValue", default_value=["bool", 0]),
|
|
"model": _node("DistributedModelName", text=["string", 0]),
|
|
})
|
|
prompt["source:1"]["inputs"].update(seed=["seed", 0], cfg=["value", 0], enabled=["enabled", 0], model=["model", 0])
|
|
graph, bindings = _run(prompt)
|
|
assert graph["source:1"]["inputs"] == {"image": "example.png", "seed": 26, "cfg": 2.75, "enabled": True, "model": "model.safetensors"}
|
|
assert set(graph) == {"source:1", bindings[0]["node_id"]}
|
|
|
|
|
|
def test_distributed_text_value_can_be_wrapped_by_primitive_string():
|
|
prompt = {
|
|
"value": _node("DistributedModelName", text="model.safetensors"),
|
|
"text": _node("PrimitiveString", value=["value", 0]),
|
|
"collector": _collector(texts=["text", 0]),
|
|
}
|
|
graph, bindings = _run(prompt)
|
|
assert graph["text"]["inputs"]["value"] == "model.safetensors"
|
|
assert graph[bindings[0]["node_id"]]["inputs"]["texts"] == ["text", 0]
|
|
assert "value" not in graph
|
|
|
|
|
|
@pytest.mark.parametrize("class_type,input_name", [
|
|
("DistributedSeed", "seed"), ("DistributedValue", "default_value"), ("DistributedModelName", "text"),
|
|
])
|
|
def test_dynamic_scalar_resolution_stays_fail_closed(class_type, input_name):
|
|
prompt = _prompt()
|
|
prompt["dynamic"] = _node("ArbitraryValue")
|
|
prompt["scalar"] = _node(class_type, **{input_name: ["dynamic", 0]})
|
|
prompt["source:1"]["inputs"]["parameter"] = ["scalar", 0]
|
|
with pytest.raises(ValueError, match="literal or supported Primitive"):
|
|
_run(prompt)
|
|
|
|
|
|
def test_namespaced_export_ids_avoid_existing_ids_even_on_pruned_nodes():
|
|
prompt = _prompt()
|
|
_, first = _run(prompt)
|
|
original_id = first[0]["node_id"]
|
|
assert ":" in original_id and not original_id.isdigit()
|
|
prompt[original_id] = _node("ExistingUtility")
|
|
_, second = _run(prompt)
|
|
assert second[0]["node_id"] not in prompt
|
|
assert second[0]["binding_id"] == first[0]["binding_id"]
|
|
prompt[second[0]["node_id"]] = _node("AnotherUtility")
|
|
graph, third = _run(prompt)
|
|
assert third[0]["node_id"] not in prompt
|
|
assert third == _run(dict(reversed(list(prompt.items()))))[1]
|
|
assert set(graph) == {"source:1", third[0]["node_id"]}
|
|
|
|
|
|
def test_long_valid_identifiers_remain_exact_and_generated_ids_fit_contract():
|
|
collector_id = "namespace/" + "c" * 190
|
|
prompt = {"source:1": _node("Source"), collector_id: _collector(images=["source:1", 0])}
|
|
prompt[collector_id]["inputs"].update(multi_job_id="j" * 200, worker_id="p" * 200)
|
|
_, bindings = _run(prompt)
|
|
assert bindings[0]["collector_id"] == collector_id
|
|
assert bindings[0]["job_id"] == "j" * 200
|
|
assert bindings[0]["participant_id"] == "p" * 200
|
|
contract.validate_bindings(bindings)
|
|
|
|
|
|
def test_deterministic_independent_copies_and_request_identity():
|
|
prompt = _prompt()
|
|
prompt["source:1"]["_meta"] = {"nested": {"values": [1, 2, 3]}}
|
|
original = copy.deepcopy(prompt)
|
|
graph, bindings = _run(prompt)
|
|
assert prompt == original
|
|
assert (graph, bindings) == _run(prompt)
|
|
assert (graph, bindings) == _run(dict(reversed(list(prompt.items()))))
|
|
other_graph, other_bindings = compile_media_prompt(prompt, "request:2")
|
|
assert other_bindings == bindings
|
|
target = bindings[0]["node_id"]
|
|
assert other_graph[target]["inputs"]["request_id"] == "request:2"
|
|
graph["source:1"]["_meta"]["nested"]["values"].append(4)
|
|
graph[target]["inputs"]["images"][0] = "changed"
|
|
bindings[0]["collector_id"] = "changed"
|
|
assert prompt == original
|
|
assert json.loads(graph[target]["inputs"]["binding_json"])["collector_id"] == "collector:1"
|
|
assert other_graph["source:1"] == original["source:1"]
|
|
|
|
|
|
@pytest.mark.parametrize("bad", [None, [], {}, {"x": None}, {"x": []}, {"x": {}},
|
|
{"x": {"class_type": None}}, {"x": {"class_type": ""}},
|
|
{"x": {"class_type": "Source", "inputs": []}},
|
|
{"x": {"class_type": "Source", "inputs": None}},
|
|
{"x": {"class_type": "Source", "inputs": {1: "literal"}}},
|
|
])
|
|
def test_rejects_malformed_prompt_or_missing_collectors(bad):
|
|
with pytest.raises(ValueError):
|
|
_run(bad)
|
|
|
|
|
|
@pytest.mark.parametrize("bad", [None, 1, True, "", " bad", "bad ", "bad\\id", "bad\nid", "x" * 201])
|
|
def test_rejects_malformed_ids_without_normalizing(bad):
|
|
prompt = _prompt()
|
|
prompt[bad] = _node("Source")
|
|
with pytest.raises(ValueError):
|
|
_run(prompt)
|
|
|
|
|
|
@pytest.mark.parametrize("bad", [None, 1, True, "", " bad", "bad ", "bad\\id", "x" * 201])
|
|
@pytest.mark.parametrize("field", ["multi_job_id", "worker_id", "request_id"])
|
|
def test_rejects_malformed_binding_and_request_identity(field, bad):
|
|
prompt = _prompt()
|
|
request_id = "request:1"
|
|
if field == "request_id":
|
|
request_id = bad
|
|
else:
|
|
prompt["collector:1"]["inputs"][field] = bad
|
|
with pytest.raises(ValueError):
|
|
compile_media_prompt(prompt, request_id)
|
|
|
|
|
|
@pytest.mark.parametrize("field", ["multi_job_id", "worker_id"])
|
|
def test_requires_prepared_literal_hidden_identity(field):
|
|
prompt = _prompt()
|
|
del prompt["collector:1"]["inputs"][field]
|
|
with pytest.raises(ValueError):
|
|
_run(prompt)
|
|
prompt["collector:1"]["inputs"][field] = ["source:1", 0]
|
|
with pytest.raises(ValueError):
|
|
_run(prompt)
|
|
|
|
|
|
@pytest.mark.parametrize("socket", SOCKETS)
|
|
@pytest.mark.parametrize("value", [None, "literal", 5, True, [], ["source:1"], ["source:1", 0, 1], ("source:1", 0)])
|
|
def test_every_present_media_socket_requires_an_api_link(socket, value):
|
|
prompt = _prompt()
|
|
prompt["collector:1"]["inputs"][socket] = value
|
|
with pytest.raises(ValueError, match="link"):
|
|
_run(prompt)
|
|
|
|
|
|
def test_each_collector_requires_at_least_one_media_socket():
|
|
prompt = _prompt()
|
|
prompt["empty"] = _collector(load_balance=True, pass_through=False)
|
|
with pytest.raises(ValueError, match="connected|linked"):
|
|
_run(prompt)
|
|
|
|
|
|
@pytest.mark.parametrize("link", [["missing", 0], ["source:1", -1], ["source:1", 0.0],
|
|
["source:1", "0"], ["source:1", True], ["source:1", False], [1, 0],
|
|
[True, 0], [None, 0], [[], 0], [{}, 0], [" source:1", 0],
|
|
])
|
|
@pytest.mark.parametrize("location", ["collector", "ordinary", "unrelated"])
|
|
def test_rejects_invalid_links_everywhere_before_pruning(location, link):
|
|
prompt = _prompt()
|
|
prompt["1"] = _node("Source") # An integer/bool source must not be stringified to this ID.
|
|
if location == "collector":
|
|
prompt["collector:1"]["inputs"]["images"] = link
|
|
elif location == "ordinary":
|
|
prompt["source:1"]["inputs"]["parameter"] = link
|
|
else:
|
|
prompt["unrelated"] = _node("Utility", value=link)
|
|
with pytest.raises(ValueError):
|
|
_run(prompt)
|
|
|
|
|
|
@pytest.mark.parametrize("class_type", [
|
|
"DistributedCollector", "DistributedEmptyImage", "UltimateSDUpscaleDistributed",
|
|
"DistributedImageListStart", "DistributedImageListEnd", "DistributedWorkerSink",
|
|
"ImageBatchDivider", "AudioBatchDivider", "CustomDistributedNode", "custombatchdivider",
|
|
])
|
|
def test_rejects_unsupported_distributed_families_even_if_unused(class_type):
|
|
prompt = _prompt()
|
|
prompt["unsupported"] = _node(class_type)
|
|
with pytest.raises(ValueError, match="unsupported"):
|
|
_run(prompt)
|
|
|
|
|
|
@pytest.mark.parametrize("class_type", ["DistributedSeed", "DistributedValue", "DistributedModelName", "PrimitiveInt", "PrimitiveString"])
|
|
def test_rejects_invalid_known_scalar_output_slots(class_type):
|
|
prompt = _prompt()
|
|
prompt["scalar"] = _node(class_type)
|
|
prompt["source:1"]["inputs"]["parameter"] = ["scalar", 1]
|
|
with pytest.raises(ValueError, match="slot"):
|
|
_run(prompt)
|
|
|
|
|
|
@pytest.mark.parametrize("socket", SOCKETS)
|
|
def test_rejects_direct_distributed_scalar_media_links(socket):
|
|
prompt = {"scalar": _node("DistributedValue", default_value="text"), "collector": _collector(**{socket: ["scalar", 0]})}
|
|
with pytest.raises(ValueError, match="scalar"):
|
|
_run(prompt)
|
|
|
|
|
|
@pytest.mark.parametrize("indirect", [False, True])
|
|
def test_rejects_collector_chains_across_different_media(indirect):
|
|
prompt = _prompt()
|
|
source = "collector:1"
|
|
if indirect:
|
|
prompt["processing"] = _node("ImageToVideo", images=[source, 0])
|
|
source = "processing"
|
|
prompt["collector:2"] = _collector(video=[source, 0])
|
|
with pytest.raises(ValueError, match="collector-to-collector"):
|
|
_run(prompt)
|
|
|
|
|
|
@pytest.mark.parametrize("mode", ["ordinary", "self", "unrelated", "primitive", "collector"])
|
|
def test_rejects_cycles_in_entire_prompt(mode):
|
|
prompt = _prompt()
|
|
if mode == "ordinary":
|
|
prompt["source:1"]["inputs"]["other"] = ["other", 0]
|
|
prompt["other"] = _node("Utility", value=["source:1", 0])
|
|
elif mode == "self":
|
|
prompt["source:1"]["inputs"]["other"] = ["source:1", 0]
|
|
elif mode == "unrelated":
|
|
prompt["unused"] = _node("Utility", value=["unused", 0])
|
|
elif mode == "primitive":
|
|
prompt["p1"] = _node("PrimitiveInt", value=["p2", 0])
|
|
prompt["p2"] = _node("PrimitiveInt", value=["p1", 0])
|
|
else:
|
|
prompt["source:1"]["inputs"]["other"] = ["collector:1", 0]
|
|
with pytest.raises(ValueError, match="cycle"):
|
|
_run(prompt)
|
|
|
|
|
|
def test_deep_acyclic_graph_does_not_depend_on_python_recursion_limit():
|
|
prompt = {"n:0": _node("Source")}
|
|
for index in range(1, 1500):
|
|
prompt[f"n:{index}"] = _node("Utility", value=[f"n:{index - 1}", 0])
|
|
prompt["collector"] = _collector(images=["n:1499", 0])
|
|
graph, bindings = _run(prompt)
|
|
assert len(graph) == 1501
|
|
assert graph[bindings[0]["node_id"]]["inputs"]["images"] == ["n:1499", 0]
|
|
|
|
|
|
def test_binding_count_respects_bridge_contract():
|
|
prompt = {"source": _node("Source")}
|
|
for index in range(contract.MAX_ARTIFACTS + 1):
|
|
prompt[f"collector:{index}"] = _collector(images=["source", 0])
|
|
with pytest.raises(ValueError, match="count"):
|
|
_run(prompt)
|
|
|
|
|
|
def test_import_and_compile_do_not_import_comfy_or_tensor_libraries():
|
|
script = '''
|
|
import builtins, importlib, pathlib, sys, types
|
|
root = pathlib.Path(sys.argv[1])
|
|
name = "isolated_media_compile"
|
|
for suffix in ("", ".api", ".api.orchestration", ".serverless_bridge"):
|
|
package = types.ModuleType(name + suffix)
|
|
package.__path__ = [str(root.joinpath(*suffix.strip(".").split(".")))] if suffix else [str(root)]
|
|
sys.modules[package.__name__] = package
|
|
original_import = builtins.__import__
|
|
def guarded_import(name, *args, **kwargs):
|
|
if name.split(".")[0] in {"comfy", "comfy_api", "nodes", "server", "torch", "numpy"}:
|
|
raise AssertionError("unexpected runtime import: " + name)
|
|
return original_import(name, *args, **kwargs)
|
|
builtins.__import__ = guarded_import
|
|
module = importlib.import_module(name + ".api.orchestration.serverless_media_prompt")
|
|
prompt = {"s": {"class_type": "Source", "inputs": {}}, "c": {"class_type": "DistributedMediaCollector", "inputs": {"texts": ["s", 0], "multi_job_id": "job", "worker_id": "worker"}}}
|
|
graph, bindings = module.compile_media_prompt(prompt, "request")
|
|
assert graph[bindings[0]["node_id"]]["class_type"] == "MediaExportV1"
|
|
'''
|
|
result = subprocess.run([sys.executable, "-c", script, str(ROOT)], text=True, capture_output=True)
|
|
assert result.returncode == 0, result.stderr
|