import importlib.util import sys import types import unittest from pathlib import Path def _load_media_sync_module(): module_path = Path(__file__).resolve().parents[2] / "api" / "orchestration" / "media_sync.py" package_name = "dist_ms_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 network_module = types.ModuleType(f"{package_name}.utils.network") network_module.build_worker_url = lambda worker, endpoint="": f"http://localhost{endpoint}" async def _fake_session(): raise RuntimeError("network calls not used in pure-function tests") network_module.get_client_session = _fake_session sys.modules[f"{package_name}.utils.network"] = network_module trace_module = types.ModuleType(f"{package_name}.utils.trace_logger") trace_module.trace_debug = lambda *_args, **_kwargs: None trace_module.trace_info = lambda *_args, **_kwargs: None sys.modules[f"{package_name}.utils.trace_logger"] = trace_module created_aiohttp_stub = False if "aiohttp" not in sys.modules: created_aiohttp_stub = True aiohttp_module = types.ModuleType("aiohttp") class _ClientTimeout: def __init__(self, total=None): pass class _FormData: def add_field(self, *args, **kwargs): pass aiohttp_module.ClientTimeout = _ClientTimeout aiohttp_module.FormData = _FormData sys.modules["aiohttp"] = aiohttp_module spec = importlib.util.spec_from_file_location( f"{package_name}.api.orchestration.media_sync", 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) if created_aiohttp_stub: sys.modules.pop("aiohttp", None) return module ms = _load_media_sync_module() # --------------------------------------------------------------------------- # convert_paths_for_platform # --------------------------------------------------------------------------- class ConvertPathsForPlatformTests(unittest.TestCase): def test_forward_slash_target_normalises_backslashes(self): obj = {"ckpt_name": "C:\\Models\\model.safetensors"} result = ms.convert_paths_for_platform(obj, "/") self.assertEqual(result["ckpt_name"], "C:/Models/model.safetensors") def test_backslash_target_normalises_forward_slashes(self): obj = {"ckpt_name": "/models/checkpoints/model.safetensors"} result = ms.convert_paths_for_platform(obj, "\\") self.assertIn("\\", result["ckpt_name"]) self.assertNotIn("/", result["ckpt_name"]) def test_relative_media_paths_always_stay_forward_slash(self): """Relative image/video/audio paths (Comfy annotated style) must not be backslash-ified.""" obj = {"image": "subfolder/my_photo.png"} result = ms.convert_paths_for_platform(obj, "\\") self.assertEqual(result["image"], "subfolder/my_photo.png") def test_relative_audio_paths_stay_forward_slash(self): obj = {"audio": "subfolder/my_track.wav"} result = ms.convert_paths_for_platform(obj, "\\") self.assertEqual(result["audio"], "subfolder/my_track.wav") def test_annotated_relative_media_path_stays_forward_slash(self): obj = {"image": "input/frame.jpg [abc123]"} result = ms.convert_paths_for_platform(obj, "\\") self.assertIn("/", result["image"]) self.assertNotIn("\\", result["image"].split("[")[0]) def test_non_filename_strings_are_untouched(self): obj = {"prompt": "a beautiful cat", "count": 5} result = ms.convert_paths_for_platform(obj, "\\") self.assertEqual(result["prompt"], "a beautiful cat") self.assertEqual(result["count"], 5) def test_url_strings_are_untouched(self): obj = {"url": "https://example.com/model.safetensors"} result = ms.convert_paths_for_platform(obj, "\\") self.assertEqual(result["url"], "https://example.com/model.safetensors") def test_invalid_separator_returns_obj_unchanged(self): obj = {"ckpt_name": "/models/model.safetensors"} result = ms.convert_paths_for_platform(obj, "|") self.assertEqual(result, obj) def test_nested_dict_is_processed_recursively(self): obj = {"node": {"ckpt_name": "C:\\Models\\model.safetensors"}} result = ms.convert_paths_for_platform(obj, "/") self.assertEqual(result["node"]["ckpt_name"], "C:/Models/model.safetensors") def test_list_items_are_processed_recursively(self): obj = [{"ckpt_name": "C:\\Models\\model.safetensors"}, "plain string"] result = ms.convert_paths_for_platform(obj, "/") self.assertEqual(result[0]["ckpt_name"], "C:/Models/model.safetensors") self.assertEqual(result[1], "plain string") def test_non_string_scalar_values_are_untouched(self): obj = {"seed": 42, "enabled": True, "ratio": 1.5} result = ms.convert_paths_for_platform(obj, "/") self.assertEqual(result["seed"], 42) self.assertTrue(result["enabled"]) def test_absolute_unix_path_to_windows(self): obj = {"lora": "/home/user/loras/my_lora.safetensors"} result = ms.convert_paths_for_platform(obj, "\\") self.assertNotIn("/", result["lora"]) def test_already_normalised_path_is_idempotent(self): obj = {"ckpt": "C:/Models/model.safetensors"} result = ms.convert_paths_for_platform(obj, "/") self.assertEqual(result["ckpt"], "C:/Models/model.safetensors") # --------------------------------------------------------------------------- # _find_media_references # --------------------------------------------------------------------------- class FindMediaReferencesTests(unittest.TestCase): def test_finds_file_input_for_load_video(self): prompt = {"1": {"class_type": "LoadVideo", "inputs": {"file": "1 - Copy.mp4"}}} refs = ms._find_media_references(prompt) self.assertIn("1 - Copy.mp4", refs) def test_strips_annotation_suffix(self): prompt = {"1": {"class_type": "LoadImage", "inputs": {"image": "photo.jpg [abc123]"}}} refs = ms._find_media_references(prompt) self.assertIn("photo.jpg", refs) self.assertFalse(any("[" in r for r in refs)) def test_normalises_backslashes_in_path(self): prompt = {"1": {"class_type": "LoadImage", "inputs": {"image": "sub\\img.png"}}} refs = ms._find_media_references(prompt) self.assertIn("sub/img.png", refs) def test_ignores_non_media_text_inputs(self): prompt = {"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "a cat"}}} refs = ms._find_media_references(prompt) self.assertEqual(refs, []) def test_ignores_node_link_values(self): """Inputs that are [node_id, slot] lists should be ignored.""" prompt = {"1": {"class_type": "Anything", "inputs": {"image": ["2", 0]}}} refs = ms._find_media_references(prompt) self.assertEqual(refs, []) def test_deduplicates_same_file_across_nodes(self): prompt = { "1": {"class_type": "LoadImage", "inputs": {"image": "cat.png"}}, "2": {"class_type": "LoadImage", "inputs": {"image": "cat.png"}}, } refs = ms._find_media_references(prompt) self.assertEqual(len(refs), 1) def test_returns_sorted_list(self): prompt = { "1": {"class_type": "LoadImage", "inputs": {"image": "z_image.png"}}, "2": {"class_type": "LoadImage", "inputs": {"image": "a_image.jpg"}}, } refs = ms._find_media_references(prompt) self.assertEqual(refs, sorted(refs)) def test_ignores_non_dict_nodes(self): prompt = {"1": "not a node dict", "2": {"class_type": "LoadImage", "inputs": {"image": "img.png"}}} refs = ms._find_media_references(prompt) self.assertIn("img.png", refs) def test_empty_prompt_returns_empty_list(self): self.assertEqual(ms._find_media_references({}), []) def test_multiple_media_types_all_found(self): prompt = { "1": {"class_type": "LoadImage", "inputs": {"image": "frame.png"}}, "2": {"class_type": "LoadVideo", "inputs": {"video": "clip.mp4"}}, "3": {"class_type": "LoadAudio", "inputs": {"audio": "track.wav"}}, } refs = ms._find_media_references(prompt) self.assertIn("frame.png", refs) self.assertIn("clip.mp4", refs) self.assertIn("track.wav", refs) class RewritePromptMediaInputsTests(unittest.TestCase): def test_rewrites_video_file_input_to_worker_path(self): prompt = { "79": {"class_type": "LoadVideo", "inputs": {"file": "1 - Copy.mp4"}}, } ms._rewrite_prompt_media_inputs(prompt, {"1 - Copy.mp4": "videos/1 - Copy.mp4"}) self.assertEqual(prompt["79"]["inputs"]["file"], "videos/1 - Copy.mp4") def test_rewrites_audio_input_and_strips_annotation_when_matching(self): prompt = { "1": {"class_type": "LoadAudio", "inputs": {"audio": "song.wav [input]"}}, } ms._rewrite_prompt_media_inputs(prompt, {"song.wav": "song.wav"}) self.assertEqual(prompt["1"]["inputs"]["audio"], "song.wav") class AnnotatedMediaUploadTests(unittest.IsolatedAsyncioTestCase): async def test_output_temp_and_input_uploads_do_not_collide(self): from unittest.mock import AsyncMock, patch class Response: status = 200 def __init__(self, data): self.data = data async def __aenter__(self): return self async def __aexit__(self, *_args): return None async def json(self): return self.data def raise_for_status(self): pass class Form: def __init__(self): self.fields = {} def add_field(self, key, value, **kwargs): self.fields[key] = value if key == 'image': self.fields['name'] = kwargs['filename'] checked = [] def post(url, *, json=None, data=None, **kwargs): if url.endswith('/check_file'): checked.append(json['filename']) return Response({'exists': False}) self.assertEqual(data.fields['type'], 'input') self.assertEqual(data.fields['image'], b'real file bytes') return Response({'name': data.fields['name'], 'subfolder': data.fields['subfolder']}) refs = ['shared/ref.png', 'shared/ref.png [output]', 'shared/ref.png [temp]'] prompt = {str(i): {'inputs': {'image': ref}} for i, ref in enumerate(refs)} self.assertEqual(ms._find_media_references(prompt), refs) with patch.object(ms, 'get_client_session', AsyncMock(return_value=types.SimpleNamespace(post=post))), \ patch.object(ms, 'aiohttp', types.SimpleNamespace(FormData=Form, ClientTimeout=lambda **kw: kw)): paths = [await ms._upload_media_to_worker({}, ref, b'real file bytes', 'fixture-hash', 'image/png') for ref in refs] expected = ['shared/ref.png', 'distributed-media/output/fixture-hash/ref.png', 'distributed-media/temp/fixture-hash/ref.png'] self.assertEqual(checked, expected) self.assertEqual(paths, [(True, path) for path in expected]) ms._rewrite_prompt_media_inputs(prompt, dict(zip(refs, expected))) self.assertEqual([node['inputs']['image'] for node in prompt.values()], expected) if __name__ == "__main__": unittest.main() # --------------------------------------------------------------------------- # fetch_worker_has_node # --------------------------------------------------------------------------- class _FakeResponse: def __init__(self, status, payload): self.status = status self._payload = payload async def json(self): if isinstance(self._payload, Exception): raise self._payload return self._payload async def __aenter__(self): return self async def __aexit__(self, *exc): return False class _FakeSession: def __init__(self, response=None, error=None): self.response = response self.error = error self.urls = [] def get(self, url, timeout=None): self.urls.append(url) if self.error: raise self.error return self.response class FetchWorkerHasNodeTests(unittest.TestCase): def _probe(self, session): import asyncio async def _get_session(): return session original = ms.get_client_session ms.get_client_session = _get_session try: return asyncio.run(ms.fetch_worker_has_node({"id": "w1"}, "DistributedWorkerSink")) finally: ms.get_client_session = original def test_registered_node_is_reported(self): session = _FakeSession(_FakeResponse(200, {"DistributedWorkerSink": {"name": "DistributedWorkerSink"}})) self.assertTrue(self._probe(session)) self.assertEqual(session.urls, ["http://localhost/object_info/DistributedWorkerSink"]) def test_unknown_node_returns_empty_object(self): # ComfyUI answers 200 {} for classes that are not registered. self.assertFalse(self._probe(_FakeSession(_FakeResponse(200, {})))) def test_http_error_counts_as_missing(self): self.assertFalse(self._probe(_FakeSession(_FakeResponse(404, None)))) def test_bad_json_counts_as_missing(self): self.assertFalse(self._probe(_FakeSession(_FakeResponse(200, ValueError("bad json"))))) def test_connection_error_counts_as_missing(self): self.assertFalse(self._probe(_FakeSession(error=OSError("unreachable"))))