260 lines
12 KiB
Python
260 lines
12 KiB
Python
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()
|