import importlib.util import sys import tempfile import types import unittest from argparse import Namespace from pathlib import Path from unittest.mock import patch def _load_process_module(module_filename: str): module_path = Path(__file__).resolve().parents[1] / "workers" / "process" / module_filename package_name = "dist_proc_testpkg" module_name = module_filename[:-3] 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 workers_pkg = types.ModuleType(f"{package_name}.workers") workers_pkg.__path__ = [str(module_path.parents[1])] sys.modules[f"{package_name}.workers"] = workers_pkg process_pkg = types.ModuleType(f"{package_name}.workers.process") process_pkg.__path__ = [] sys.modules[f"{package_name}.workers.process"] = process_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 process_module = types.ModuleType(f"{package_name}.utils.process") process_module.get_python_executable = lambda: "/usr/bin/test-python" process_module.is_process_alive = lambda _pid: False process_module.terminate_process = lambda *_args, **_kwargs: None sys.modules[f"{package_name}.utils.process"] = process_module config_module = types.ModuleType(f"{package_name}.utils.config") config_module.load_config = lambda: {"workers": [], "settings": {"stop_workers_on_master_exit": False}} config_module.save_config = lambda _config: None sys.modules[f"{package_name}.utils.config"] = config_module constants_module = types.ModuleType(f"{package_name}.utils.constants") constants_module.PROCESS_TERMINATION_TIMEOUT = 1 constants_module.PROCESS_WAIT_TIMEOUT = 1 constants_module.WORKER_CHECK_INTERVAL = 0.01 sys.modules[f"{package_name}.utils.constants"] = constants_module network_module = types.ModuleType(f"{package_name}.utils.network") network_module.get_server_port = lambda: 8189 sys.modules[f"{package_name}.utils.network"] = network_module spec = importlib.util.spec_from_file_location( f"{package_name}.workers.process.{module_name}", 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 root_discovery_module = _load_process_module("root_discovery.py") launch_builder_module = _load_process_module("launch_builder.py") class ComfyRootDiscoveryTests(unittest.TestCase): def test_prefers_loaded_comfyui_module_path(self): discovery = root_discovery_module.ComfyRootDiscovery() server_module = types.SimpleNamespace(__file__="/opt/ComfyUI/server.py") def fake_exists(path): return path == "/opt/ComfyUI/main.py" with patch.dict(sys.modules, {"server": server_module}, clear=False), \ patch.object(root_discovery_module.os.path, "exists", side_effect=fake_exists), \ patch.dict(root_discovery_module.os.environ, {}, clear=True): self.assertEqual(discovery.find_comfy_root(), "/opt/ComfyUI") class LaunchCommandBuilderTests(unittest.TestCase): def build_command(self, root, worker=None, runtime=None): builder = launch_builder_module.LaunchCommandBuilder() worker = worker or {"id": "worker-a", "port": 8190} runtime = runtime if runtime is not None else Namespace(database_url=None) with patch.object(builder, "_get_runtime_args", return_value=runtime): return builder.build_launch_command(worker, str(root)) def test_database_is_unique_and_stable_by_worker_id(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) (root / "main.py").touch() first = self.build_command(root, {"id": "worker-a", "port": 8190}) renamed = self.build_command(root, { "id": "worker-a", "port": 9001, "name": "Renamed", "cuda_device": 3, }) second = self.build_command(root, {"id": "worker-b", "port": 8191}) database = first[first.index("--database-url") + 1] self.assertEqual(database, renamed[renamed.index("--database-url") + 1]) self.assertNotEqual(database, second[second.index("--database-url") + 1]) self.assertTrue(database.startswith("sqlite:///" + root.as_posix() + "/")) self.assertTrue(Path(database.removeprefix("sqlite:///")).parent.is_dir()) self.assertNotIn("--disable-assets", first) def test_database_uses_effective_user_or_base_directory(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) (root / "main.py").touch() for extra_args, expected in [ (f'--user-directory "{root / "custom user"}"', root / "custom user"), (f'--base-directory "{root / "custom base"}"', root / "custom base" / "user"), ]: with self.subTest(extra_args=extra_args): cmd = self.build_command(root, { "id": "worker-a", "port": 8190, "extra_args": extra_args, }) database = cmd[cmd.index("--database-url") + 1] self.assertTrue(database.startswith("sqlite:///" + expected.as_posix() + "/")) cmd = self.build_command(root, runtime=Namespace( database_url="sqlite:///master.db", user_directory=str(root / "runtime user"), )) database = cmd[cmd.index("--database-url") + 1] self.assertTrue(database.startswith("sqlite:///" + (root / "runtime user").as_posix() + "/")) self.assertNotEqual(database, "sqlite:///master.db") def test_explicit_database_and_disabled_assets_are_preserved(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) (root / "main.py").touch() for extra in ["--database-url sqlite:///explicit.db", "--database-url=sqlite:///explicit.db", "--disable-assets"]: with self.subTest(extra=extra): cmd = self.build_command(root, {"id": "worker-a", "port": 8190, "extra_args": extra}) self.assertEqual(sum(arg.split("=")[0] == "--database-url" for arg in cmd), 0 if extra == "--disable-assets" else 1) self.assertFalse((root / "user").exists()) def test_older_comfyui_without_database_flag_keeps_existing_command(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) (root / "main.py").touch() cmd = self.build_command(root, runtime=Namespace()) self.assertNotIn("--database-url", cmd) self.assertFalse((root / "user").exists()) def test_worker_id_cannot_escape_database_directory(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) (root / "main.py").touch() databases = [] for worker_id in ["../outside", "a/b", "a_b"]: cmd = self.build_command(root, {"id": worker_id, "port": 8190}) path = Path(cmd[cmd.index("--database-url") + 1].removeprefix("sqlite:///")) self.assertTrue(path.is_relative_to(root / "user")) databases.append(path) self.assertEqual(len(set(databases)), 3) def test_extra_args_cannot_silently_override_the_configured_port(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) (root / "main.py").touch() for extra in ["--port 8189", "--port=8189"]: with self.subTest(extra=extra), self.assertRaisesRegex(ValueError, "configured worker port"): self.build_command(root, {"id": "worker-a", "port": 8190, "extra_args": extra}) def test_inherits_runtime_layout_args_for_desktop(self): builder = launch_builder_module.LaunchCommandBuilder() runtime_args = Namespace( listen="127.0.0.1", base_directory="C:/Users/test/ComfyUI", temp_directory=None, input_directory="C:/Users/test/ComfyUI/input", output_directory="C:/Users/test/ComfyUI/output", user_directory="C:/Users/test/ComfyUI/user", front_end_root="C:/Program Files/ComfyUI/web_custom_versions/desktop_app", extra_model_paths_config=[["C:/Users/test/AppData/Roaming/ComfyUI/extra_models_config.yaml"]], enable_manager=True, disable_manager_ui=False, enable_manager_legacy_ui=False, windows_standalone_build=True, log_stdout=True, verbose="INFO", enable_cors_header="*", ) comfy_module = types.ModuleType("comfy") comfy_cli_args = types.ModuleType("comfy.cli_args") comfy_cli_args.args = runtime_args worker_config = { "port": 9001, "extra_args": "--preview-method auto", } def fake_exists(path): return path == "/desktop/ComfyUI/main.py" with patch.dict( sys.modules, {"comfy": comfy_module, "comfy.cli_args": comfy_cli_args}, clear=False, ), patch.object(launch_builder_module.os.path, "exists", side_effect=fake_exists): cmd = builder.build_launch_command(worker_config, "/desktop/ComfyUI") self.assertEqual(cmd[:2], ["/usr/bin/test-python", "/desktop/ComfyUI/main.py"]) self.assertIn("--listen", cmd) self.assertIn("127.0.0.1", cmd) self.assertIn("--base-directory", cmd) self.assertIn("C:/Users/test/ComfyUI", cmd) self.assertIn("--input-directory", cmd) self.assertIn("--output-directory", cmd) self.assertIn("--user-directory", cmd) self.assertIn("--front-end-root", cmd) self.assertIn("--extra-model-paths-config", cmd) self.assertIn("C:/Users/test/AppData/Roaming/ComfyUI/extra_models_config.yaml", cmd) self.assertIn("--enable-manager", cmd) self.assertIn("--windows-standalone-build", cmd) self.assertIn("--log-stdout", cmd) self.assertIn("--disable-auto-launch", cmd) self.assertIn("--enable-cors-header", cmd) self.assertIn("*", cmd) self.assertIn("--port", cmd) self.assertIn("9001", cmd) self.assertNotIn("--auto-launch", cmd) class ProcessLaunchPortTests(unittest.TestCase): def test_conflict_is_rejected_before_building_or_spawning(self): lifecycle_module = _load_process_module("lifecycle.py") manager = types.SimpleNamespace(find_comfy_root=lambda: "/ComfyUI", processes={}) with patch.object(lifecycle_module.subprocess, "Popen") as spawn: with self.assertRaisesRegex(ValueError, "master port 8189"): lifecycle_module.ProcessLifecycle(manager).launch_worker({"id": "manual", "port": 8189}) spawn.assert_not_called() self.assertEqual(manager.processes, {}) def test_available_port_reaches_existing_launch_path(self): lifecycle_module = _load_process_module("lifecycle.py") with tempfile.TemporaryDirectory() as directory: manager = types.SimpleNamespace( find_comfy_root=lambda: directory, build_launch_command=lambda _worker, _root: ["python", "main.py", "--port", "8190"], processes={}, save_processes=lambda: None, ) with patch.object(lifecycle_module, "validate_worker_port") as validate, \ patch.object(lifecycle_module.subprocess, "Popen", return_value=types.SimpleNamespace(pid=1234)) as spawn: pid = lifecycle_module.ProcessLifecycle(manager).launch_worker({"id": "manual", "port": 8190}) validate.assert_called_once_with({"id": "manual", "port": 8190}, 8189, []) spawn.assert_called_once() self.assertEqual(pid, 1234) self.assertEqual(manager.processes["manual"]["pid"], 1234) if __name__ == "__main__": unittest.main()