263 lines
11 KiB
Python
263 lines
11 KiB
Python
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_image_input(self):
|
|
prompt = {"1": {"class_type": "LoadImage", "inputs": {"image": "photo.png"}}}
|
|
refs = ms._find_media_references(prompt)
|
|
self.assertIn("photo.png", refs)
|
|
|
|
def test_finds_video_input(self):
|
|
prompt = {"1": {"class_type": "LoadVideo", "inputs": {"video": "clip.mp4"}}}
|
|
refs = ms._find_media_references(prompt)
|
|
self.assertIn("clip.mp4", refs)
|
|
|
|
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_finds_audio_input(self):
|
|
prompt = {"1": {"class_type": "LoadAudio", "inputs": {"audio": "track.wav"}}}
|
|
refs = ms._find_media_references(prompt)
|
|
self.assertIn("track.wav", 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")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|