import importlib.util import json import sys import types import unittest from pathlib import Path def _load_prompt_transform_module(): module_path = Path(__file__).resolve().parents[1] / "api" / "orchestration" / "prompt_transform.py" package_name = "dist_pt_testpkg" for mod_name in list(sys.modules): if mod_name == package_name or mod_name.startswith(f"{package_name}."): del sys.modules[mod_name] root_pkg = types.ModuleType(package_name) root_pkg.__path__ = [] sys.modules[package_name] = root_pkg api_pkg = types.ModuleType(f"{package_name}.api") api_pkg.__path__ = [] sys.modules[f"{package_name}.api"] = api_pkg orch_pkg = types.ModuleType(f"{package_name}.api.orchestration") orch_pkg.__path__ = [] sys.modules[f"{package_name}.api.orchestration"] = orch_pkg utils_pkg = types.ModuleType(f"{package_name}.utils") utils_pkg.__path__ = [] sys.modules[f"{package_name}.utils"] = utils_pkg logging_module = types.ModuleType(f"{package_name}.utils.logging") logging_module.debug_log = lambda *_args, **_kwargs: None logging_module.log = lambda *_args, **_kwargs: None sys.modules[f"{package_name}.utils.logging"] = logging_module spec = importlib.util.spec_from_file_location( f"{package_name}.api.orchestration.prompt_transform", module_path, ) module = importlib.util.module_from_spec(spec) assert spec is not None and spec.loader is not None spec.loader.exec_module(module) return module pt = _load_prompt_transform_module() # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _linear_prompt(): """1 → 2 → 3 → 4(DistributedCollector) → 5(SaveImage)""" return { "1": {"class_type": "CheckpointLoaderSimple", "inputs": {}}, "2": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["1", 1]}}, "3": {"class_type": "KSampler", "inputs": {"model": ["1", 0], "positive": ["2", 0]}}, "4": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}}, "5": {"class_type": "SaveImage", "inputs": {"images": ["4", 0]}}, } def _collector_only_prompt(): """1(Checkpoint) → 2(DistributedCollector) [no downstream from 2]""" return { "1": {"class_type": "CheckpointLoaderSimple", "inputs": {}}, "2": {"class_type": "DistributedCollector", "inputs": {"images": ["1", 0]}}, } def _delegate_prompt(): """1 → 2 → 3(DistributedCollector) → 4(SaveImage)""" return { "1": {"class_type": "CheckpointLoaderSimple", "inputs": {}}, "2": {"class_type": "KSampler", "inputs": {"model": ["1", 0]}}, "3": {"class_type": "DistributedCollector", "inputs": {"images": ["2", 0]}}, "4": {"class_type": "SaveImage", "inputs": {"images": ["3", 0]}}, } def _apply(prompt, participant_id, enabled_worker_ids=None, delegate_master=False): if enabled_worker_ids is None: enabled_worker_ids = ["worker-a", "worker-b"] idx = pt.PromptIndex(prompt) job_id_map = pt.generate_job_id_map(idx, "run") return pt.apply_participant_overrides( prompt, participant_id=participant_id, enabled_worker_ids=enabled_worker_ids, job_id_map=job_id_map, master_url="http://master.example.com", delegate_master=delegate_master, prompt_index=idx, ) # --------------------------------------------------------------------------- # PromptIndex # --------------------------------------------------------------------------- class PromptIndexTests(unittest.TestCase): def test_nodes_by_class_groups_correctly(self): prompt = { "1": {"class_type": "CheckpointLoaderSimple", "inputs": {}}, "2": {"class_type": "DistributedCollector", "inputs": {}}, "3": {"class_type": "DistributedCollector", "inputs": {}}, } idx = pt.PromptIndex(prompt) self.assertCountEqual(idx.nodes_for_class("DistributedCollector"), ["2", "3"]) self.assertEqual(idx.nodes_for_class("CheckpointLoaderSimple"), ["1"]) def test_nodes_for_class_unknown_returns_empty(self): idx = pt.PromptIndex({"1": {"class_type": "KSampler", "inputs": {}}}) self.assertEqual(idx.nodes_for_class("Nonexistent"), []) def test_nodes_without_class_type_are_indexed_under_none(self): prompt = {"1": {"inputs": {}}} idx = pt.PromptIndex(prompt) # Should not raise; nodes_for_class with None key or missing class_type self.assertEqual(idx.nodes_for_class("KSampler"), []) def test_copy_prompt_is_a_deep_copy(self): prompt = {"1": {"class_type": "KSampler", "inputs": {"seed": 42}}} idx = pt.PromptIndex(prompt) copy = idx.copy_prompt() copy["1"]["inputs"]["seed"] = 999 self.assertEqual(prompt["1"]["inputs"]["seed"], 42) def test_has_upstream_direct_connection(self): """Node 4 reads directly from node 3 (KSampler).""" idx = pt.PromptIndex(_linear_prompt()) self.assertTrue(idx.has_upstream("4", "KSampler")) def test_has_upstream_transitive_connection(self): """Node 4 → 3 → 2 → 1 (CheckpointLoaderSimple).""" idx = pt.PromptIndex(_linear_prompt()) self.assertTrue(idx.has_upstream("4", "CheckpointLoaderSimple")) def test_has_upstream_returns_false_when_no_path(self): idx = pt.PromptIndex(_linear_prompt()) # CheckpointLoaderSimple has no upstream nodes self.assertFalse(idx.has_upstream("1", "DistributedCollector")) def test_has_upstream_result_is_cached(self): idx = pt.PromptIndex(_linear_prompt()) r1 = idx.has_upstream("4", "KSampler") r2 = idx.has_upstream("4", "KSampler") self.assertEqual(r1, r2) self.assertIn(("4", "KSampler"), idx._upstream_cache) def test_has_upstream_does_not_infinite_loop_on_cycle(self): """Cyclic references in inputs should not cause infinite recursion.""" prompt = { "1": {"class_type": "A", "inputs": {"x": ["2", 0]}}, "2": {"class_type": "B", "inputs": {"x": ["1", 0]}}, } idx = pt.PromptIndex(prompt) # Should terminate without error result = idx.has_upstream("1", "NonExistent") self.assertFalse(result) # --------------------------------------------------------------------------- # find_nodes_by_class # --------------------------------------------------------------------------- class FindNodesByClassTests(unittest.TestCase): def test_finds_matching_nodes(self): prompt = { "1": {"class_type": "KSampler", "inputs": {}}, "2": {"class_type": "DistributedCollector", "inputs": {}}, } result = pt.find_nodes_by_class(prompt, "KSampler") self.assertEqual(result, ["1"]) def test_returns_empty_when_no_match(self): prompt = {"1": {"class_type": "KSampler", "inputs": {}}} self.assertEqual(pt.find_nodes_by_class(prompt, "DistributedCollector"), []) def test_skips_non_dict_nodes(self): prompt = {"1": "not a dict", "2": {"class_type": "KSampler", "inputs": {}}} result = pt.find_nodes_by_class(prompt, "KSampler") self.assertEqual(result, ["2"]) # --------------------------------------------------------------------------- # prune_prompt_for_worker # --------------------------------------------------------------------------- class PrunePromptForWorkerTests(unittest.TestCase): def test_no_distributed_nodes_returns_prompt_unchanged(self): prompt = { "1": {"class_type": "CheckpointLoaderSimple", "inputs": {}}, "2": {"class_type": "SaveImage", "inputs": {"images": ["1", 0]}}, } result = pt.prune_prompt_for_worker(prompt) self.assertCountEqual(result.keys(), ["1", "2"]) def test_keeps_collector_and_upstream(self): prompt = _linear_prompt() result = pt.prune_prompt_for_worker(prompt) for node_id in ("1", "2", "3", "4"): self.assertIn(node_id, result) def test_removes_downstream_of_collector(self): prompt = _linear_prompt() result = pt.prune_prompt_for_worker(prompt) self.assertNotIn("5", result) def test_injects_preview_image_when_downstream_exists(self): prompt = _linear_prompt() result = pt.prune_prompt_for_worker(prompt) preview_nodes = [n for n in result.values() if n.get("class_type") == "PreviewImage"] self.assertEqual(len(preview_nodes), 1) self.assertEqual(preview_nodes[0]["inputs"]["images"], ["4", 0]) def test_no_preview_image_when_no_downstream(self): result = pt.prune_prompt_for_worker(_collector_only_prompt()) preview_nodes = [n for n in result.values() if n.get("class_type") == "PreviewImage"] self.assertEqual(len(preview_nodes), 0) def test_unrelated_nodes_are_pruned(self): prompt = { "1": {"class_type": "DistributedCollector", "inputs": {}}, "2": {"class_type": "UnrelatedNode", "inputs": {}}, # no connection to 1 } result = pt.prune_prompt_for_worker(prompt) self.assertIn("1", result) self.assertNotIn("2", result) def test_result_is_a_copy_not_same_object(self): prompt = _linear_prompt() result = pt.prune_prompt_for_worker(prompt) # Mutating the result should not affect the original original_keys = set(prompt.keys()) result["NEW"] = {"class_type": "Test", "inputs": {}} self.assertEqual(set(prompt.keys()), original_keys) def test_upscale_node_is_treated_as_distributed(self): prompt = { "1": {"class_type": "KSampler", "inputs": {}}, "2": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"image": ["1", 0]}}, "3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0]}}, } result = pt.prune_prompt_for_worker(prompt) self.assertIn("1", result) self.assertIn("2", result) self.assertNotIn("3", result) # --------------------------------------------------------------------------- # prepare_delegate_master_prompt # --------------------------------------------------------------------------- class PrepareDelegateMasterPromptTests(unittest.TestCase): def test_keeps_collector_and_downstream(self): prompt = _delegate_prompt() result = pt.prepare_delegate_master_prompt(prompt, ["3"]) self.assertIn("3", result) self.assertIn("4", result) self.assertNotIn("1", result) self.assertNotIn("2", result) def test_removes_dangling_upstream_refs(self): """Collector must not retain dangling refs to pruned upstream nodes.""" prompt = _delegate_prompt() result = pt.prepare_delegate_master_prompt(prompt, ["3"]) collector_inputs = result["3"].get("inputs", {}) # Original "images" pointed at node 2, which is pruned. # It should now point at a newly injected placeholder node. self.assertIn("images", collector_inputs) source_id = str(collector_inputs["images"][0]) self.assertNotEqual(source_id, "2") self.assertIn(source_id, result) self.assertEqual(result[source_id].get("class_type"), "DistributedEmptyImage") def test_injects_empty_image_placeholder(self): prompt = _delegate_prompt() result = pt.prepare_delegate_master_prompt(prompt, ["3"]) empty_nodes = [(nid, n) for nid, n in result.items() if n.get("class_type") == "DistributedEmptyImage"] self.assertEqual(len(empty_nodes), 1) placeholder_id = empty_nodes[0][0] self.assertEqual(result["3"]["inputs"]["images"], [placeholder_id, 0]) def test_one_placeholder_per_collector(self): """Two collectors → two placeholders.""" prompt = { "1": {"class_type": "DistributedCollector", "inputs": {}}, "2": {"class_type": "DistributedCollector", "inputs": {}}, "3": {"class_type": "SaveImage", "inputs": {"images": ["1", 0]}}, } result = pt.prepare_delegate_master_prompt(prompt, ["1", "2"]) empty_nodes = [n for n in result.values() if n.get("class_type") == "DistributedEmptyImage"] self.assertEqual(len(empty_nodes), 2) def test_result_is_independent_copy(self): prompt = _delegate_prompt() result = pt.prepare_delegate_master_prompt(prompt, ["3"]) result["3"]["inputs"]["NEW"] = "injected" # Original should be untouched self.assertNotIn("NEW", prompt["3"].get("inputs", {})) # --------------------------------------------------------------------------- # generate_job_id_map # --------------------------------------------------------------------------- class GenerateJobIdMapTests(unittest.TestCase): def test_maps_collector_nodes(self): prompt = { "1": {"class_type": "DistributedCollector", "inputs": {}}, "2": {"class_type": "KSampler", "inputs": {}}, } idx = pt.PromptIndex(prompt) job_map = pt.generate_job_id_map(idx, "prefix") self.assertEqual(job_map["1"], "prefix_1") self.assertNotIn("2", job_map) def test_maps_upscale_nodes(self): prompt = { "5": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {}}, } idx = pt.PromptIndex(prompt) job_map = pt.generate_job_id_map(idx, "run") self.assertEqual(job_map["5"], "run_5") def test_empty_prompt_returns_empty_map(self): idx = pt.PromptIndex({}) self.assertEqual(pt.generate_job_id_map(idx, "prefix"), {}) def test_stable_ids_across_calls(self): prompt = {"1": {"class_type": "DistributedCollector", "inputs": {}}} idx = pt.PromptIndex(prompt) m1 = pt.generate_job_id_map(idx, "run") m2 = pt.generate_job_id_map(idx, "run") self.assertEqual(m1, m2) # --------------------------------------------------------------------------- # apply_participant_overrides – DistributedCollector # --------------------------------------------------------------------------- class ApplyOverridesCollectorTests(unittest.TestCase): def _collector_prompt(self): return {"1": {"class_type": "DistributedCollector", "inputs": {}}} def test_worker_sets_is_worker_true(self): result = _apply(self._collector_prompt(), "worker-a") self.assertTrue(result["1"]["inputs"]["is_worker"]) def test_worker_sets_master_url(self): result = _apply(self._collector_prompt(), "worker-a") self.assertEqual(result["1"]["inputs"]["master_url"], "http://master.example.com") def test_worker_sets_worker_id(self): result = _apply(self._collector_prompt(), "worker-a") self.assertEqual(result["1"]["inputs"]["worker_id"], "worker-a") def test_worker_sets_delegate_only_false(self): result = _apply(self._collector_prompt(), "worker-a") self.assertFalse(result["1"]["inputs"]["delegate_only"]) def test_master_sets_is_worker_false(self): result = _apply(self._collector_prompt(), "master") self.assertFalse(result["1"]["inputs"]["is_worker"]) def test_master_clears_stale_master_url(self): prompt = {"1": {"class_type": "DistributedCollector", "inputs": {"master_url": "stale"}}} result = _apply(prompt, "master") self.assertNotIn("master_url", result["1"]["inputs"]) def test_master_clears_stale_worker_id(self): prompt = {"1": {"class_type": "DistributedCollector", "inputs": {"worker_id": "stale"}}} result = _apply(prompt, "master") self.assertNotIn("worker_id", result["1"]["inputs"]) def test_master_with_delegate_master_sets_delegate_only_true(self): result = _apply(self._collector_prompt(), "master", delegate_master=True) self.assertTrue(result["1"]["inputs"]["delegate_only"]) def test_master_without_delegate_master_sets_delegate_only_false(self): result = _apply(self._collector_prompt(), "master", delegate_master=False) self.assertFalse(result["1"]["inputs"]["delegate_only"]) def test_enabled_worker_ids_serialized_as_json(self): enabled = ["worker-a", "worker-b"] result = _apply(self._collector_prompt(), "master", enabled_worker_ids=enabled) self.assertEqual(result["1"]["inputs"]["enabled_worker_ids"], json.dumps(enabled)) def test_multi_job_id_is_set_from_job_map(self): prompt = {"1": {"class_type": "DistributedCollector", "inputs": {}}} idx = pt.PromptIndex(prompt) job_id_map = {"1": "run_abc_1"} result = pt.apply_participant_overrides( prompt, participant_id="worker-a", enabled_worker_ids=["worker-a"], job_id_map=job_id_map, master_url="http://master", delegate_master=False, prompt_index=idx, ) self.assertEqual(result["1"]["inputs"]["multi_job_id"], "run_abc_1") # --------------------------------------------------------------------------- # apply_participant_overrides – DistributedSeed # --------------------------------------------------------------------------- class ApplyOverridesSeedTests(unittest.TestCase): def _seed_prompt(self): return {"1": {"class_type": "DistributedSeed", "inputs": {}}} def test_worker_sets_is_worker_true(self): result = _apply(self._seed_prompt(), "worker-a") self.assertTrue(result["1"]["inputs"]["is_worker"]) def test_worker_id_reflects_index_in_enabled_list(self): result = _apply(self._seed_prompt(), "worker-b", enabled_worker_ids=["worker-a", "worker-b"]) self.assertEqual(result["1"]["inputs"]["worker_id"], "worker_1") def test_master_sets_is_worker_false(self): result = _apply(self._seed_prompt(), "master") self.assertFalse(result["1"]["inputs"]["is_worker"]) def test_master_sets_empty_worker_id(self): result = _apply(self._seed_prompt(), "master") self.assertEqual(result["1"]["inputs"]["worker_id"], "") # --------------------------------------------------------------------------- # apply_participant_overrides – UltimateSDUpscaleDistributed # --------------------------------------------------------------------------- class ApplyOverridesUpscaleTests(unittest.TestCase): def _upscale_prompt(self): return {"1": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {}}} def test_worker_sets_is_worker_true(self): result = _apply(self._upscale_prompt(), "worker-a") self.assertTrue(result["1"]["inputs"]["is_worker"]) def test_worker_sets_master_url_and_worker_id(self): result = _apply(self._upscale_prompt(), "worker-a") self.assertEqual(result["1"]["inputs"]["master_url"], "http://master.example.com") self.assertEqual(result["1"]["inputs"]["worker_id"], "worker-a") def test_master_clears_master_url_and_worker_id(self): prompt = {"1": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"master_url": "x", "worker_id": "y"}}} result = _apply(prompt, "master") self.assertNotIn("master_url", result["1"]["inputs"]) self.assertNotIn("worker_id", result["1"]["inputs"]) def test_collector_downstream_of_upscale_gets_pass_through(self): """A DistributedCollector that is downstream of UltimateSDUpscaleDistributed → pass_through=True.""" prompt = { "1": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {}}, "2": {"class_type": "DistributedCollector", "inputs": {"images": ["1", 0]}}, } result = _apply(prompt, "worker-a", enabled_worker_ids=["worker-a"]) self.assertTrue(result["2"]["inputs"].get("pass_through")) # --------------------------------------------------------------------------- # apply_participant_overrides – DistributedValue # --------------------------------------------------------------------------- class ApplyOverridesValueTests(unittest.TestCase): def _value_prompt(self): return {"1": {"class_type": "DistributedValue", "inputs": {}}} def test_worker_sets_is_worker_true(self): result = _apply(self._value_prompt(), "worker-a") self.assertTrue(result["1"]["inputs"]["is_worker"]) def test_worker_id_reflects_index_in_enabled_list(self): result = _apply(self._value_prompt(), "worker-b", enabled_worker_ids=["worker-a", "worker-b"]) self.assertEqual(result["1"]["inputs"]["worker_id"], "worker_1") def test_master_sets_is_worker_false(self): result = _apply(self._value_prompt(), "master") self.assertFalse(result["1"]["inputs"]["is_worker"]) def test_master_sets_empty_worker_id(self): result = _apply(self._value_prompt(), "master") self.assertEqual(result["1"]["inputs"]["worker_id"], "") if __name__ == "__main__": unittest.main()