import importlib.util import os import sys import tempfile import types import unittest from argparse import Namespace from pathlib import Path from unittest.mock import patch from sqlalchemy import create_engine from sqlalchemy.engine import URL, make_url 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_relative_database_directories_use_worker_cwd(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) / "ComfyUI" root.mkdir() (root / "main.py").touch() launcher = Path(directory) / "launcher" launcher.mkdir() cases = [ (Namespace(database_url=None, user_directory="profiles"), "", root / "profiles"), (Namespace(database_url=None, base_directory="data"), "", root / "data/user"), (Namespace(database_url=None), "--user-directory=profiles", root / "profiles"), (Namespace(database_url=None), "--base-directory data", root / "data/user"), ] for runtime, extra, expected in cases: with self.subTest(runtime=runtime, extra=extra), \ patch.object(os, "getcwd", return_value=str(launcher)): cmd = self.build_command(root, { "id": "worker-a", "port": 8190, "extra_args": extra, }, runtime) database = make_url(cmd[cmd.index("--database-url") + 1]).database self.assertEqual(Path(database).parent, expected / "distributed/workers") self.assertFalse((launcher / "profiles").exists()) self.assertFalse((launcher / "data").exists()) def test_database_url_round_trips_special_characters_without_collisions(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) (root / "main.py").touch() for name in ["user?profile", "user%3Fprofile"]: with self.subTest(name=name): user = root / name runtime = Namespace(database_url=None, user_directory=str(user)) # SQLAlchemy 2.0 cannot round-trip '?' in a filename. Fail # clearly on those versions rather than silently sharing a DB. sample = (user / "test.db").as_posix() serialized = URL.create("sqlite", database=sample).render_as_string() if make_url(serialized).database != sample: with self.assertRaisesRegex(ValueError, "SQLAlchemy.*database path"): self.build_command(root, runtime=runtime) self.assertFalse(user.exists()) continue databases = [] for worker_id in ["worker-a", "worker-b"]: cmd = self.build_command(root, {"id": worker_id, "port": 8190}, runtime) url = cmd[cmd.index("--database-url") + 1] database = make_url(url).database self.assertEqual(Path(database).parent, user / "distributed/workers") engine = create_engine(url) try: with engine.connect() as connection: actual = connection.exec_driver_sql("PRAGMA database_list").one()[2] self.assertEqual(Path(actual), Path(database)) finally: engine.dispose() databases.append(database) self.assertNotEqual(*databases) 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_launch_before_master_listens_uses_configured_master_port(self): lifecycle_module = _load_process_module("lifecycle.py") cli_args = types.ModuleType("comfy.cli_args") cli_args.args = Namespace(port=8189) for missing in [AttributeError("PromptServer has no port"), None]: with self.subTest(missing=missing), tempfile.TemporaryDirectory() as directory: manager = types.SimpleNamespace( find_comfy_root=lambda: directory, build_launch_command=lambda _worker, _root: ["python", "main.py"], processes={}, save_processes=lambda: None, ) with patch.dict(sys.modules, {"comfy": types.ModuleType("comfy"), "comfy.cli_args": cli_args}), \ patch.object(lifecycle_module, "get_server_port", side_effect=missing if isinstance(missing, Exception) else None, return_value=None), \ 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) def test_pre_listen_launch_still_rejects_the_configured_master_port(self): lifecycle_module = _load_process_module("lifecycle.py") cli_args = types.ModuleType("comfy.cli_args") cli_args.args = Namespace(port=8189) manager = types.SimpleNamespace(find_comfy_root=lambda: "/ComfyUI", processes={}) with patch.dict(sys.modules, {"comfy": types.ModuleType("comfy"), "comfy.cli_args": cli_args}), \ patch.object(lifecycle_module, "get_server_port", side_effect=AttributeError("port")), \ 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() 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()