112 lines
3.7 KiB
Python
112 lines
3.7 KiB
Python
"""Load the plugin package under test with lightweight ComfyUI stubs."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import sys
|
|
import types
|
|
from pathlib import Path
|
|
from types import ModuleType
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[2]
|
|
PLUGIN_INIT_PATH = REPO_ROOT / "__init__.py"
|
|
|
|
|
|
class RecordingProgressBar:
|
|
"""Capture progress updates emitted through the ComfyUI progress boundary."""
|
|
|
|
instances: list["RecordingProgressBar"] = []
|
|
|
|
def __init__(self, total: int):
|
|
"""Store the requested total and record this instance globally."""
|
|
self.total = total
|
|
self.calls: list[tuple[int, int, tuple[str, object, int] | None]] = []
|
|
self.__class__.instances.append(self)
|
|
|
|
def update_absolute(
|
|
self,
|
|
value: int,
|
|
total: int,
|
|
preview: tuple[str, object, int] | None = None,
|
|
) -> None:
|
|
"""Record each progress update for later assertions."""
|
|
self.calls.append((value, total, preview))
|
|
|
|
@classmethod
|
|
def reset(cls) -> None:
|
|
"""Clear recorded instances between tests."""
|
|
cls.instances = []
|
|
|
|
|
|
def install_comfy_stubs(
|
|
progress_bar_class: type[RecordingProgressBar] | None = None,
|
|
) -> None:
|
|
"""Install minimal Comfy and Comfy API stubs needed by plugin imports."""
|
|
progress_bar = progress_bar_class or RecordingProgressBar
|
|
comfy_module = types.ModuleType("comfy")
|
|
comfy_utils_module = types.ModuleType("comfy.utils")
|
|
comfy_utils_module.ProgressBar = progress_bar
|
|
comfy_module.utils = comfy_utils_module
|
|
comfy_api_module = sys.modules.get("comfy_api") or types.ModuleType("comfy_api")
|
|
comfy_api_latest_module = sys.modules.get("comfy_api.latest") or types.ModuleType(
|
|
"comfy_api.latest"
|
|
)
|
|
|
|
existing_input_impl = getattr(comfy_api_latest_module, "InputImpl", None)
|
|
if existing_input_impl is not None and hasattr(
|
|
existing_input_impl, "VideoFromFile"
|
|
):
|
|
video_from_file = existing_input_impl.VideoFromFile
|
|
else:
|
|
|
|
class VideoFromFile:
|
|
"""Store one file-backed video stream source for test-only Comfy stubs."""
|
|
|
|
def __init__(self, file):
|
|
"""Store the provided stream source for later retrieval."""
|
|
self._file = file
|
|
|
|
def get_stream_source(self):
|
|
"""Return the stored stream source unchanged."""
|
|
return self._file
|
|
|
|
video_from_file = VideoFromFile
|
|
|
|
comfy_api_latest_module.InputImpl = types.SimpleNamespace(
|
|
VideoFromFile=video_from_file
|
|
)
|
|
|
|
sys.modules["comfy"] = comfy_module
|
|
sys.modules["comfy.utils"] = comfy_utils_module
|
|
sys.modules["comfy_api"] = comfy_api_module
|
|
sys.modules["comfy_api.latest"] = comfy_api_latest_module
|
|
|
|
|
|
def unload_package(package_name: str) -> None:
|
|
"""Remove a previously loaded plugin package alias from `sys.modules`."""
|
|
for module_name in list(sys.modules):
|
|
if module_name == package_name or module_name.startswith(f"{package_name}."):
|
|
del sys.modules[module_name]
|
|
|
|
|
|
def load_plugin_package(
|
|
package_name: str = "comfyui_mutiny_under_test",
|
|
progress_bar_class: type[RecordingProgressBar] | None = None,
|
|
) -> ModuleType:
|
|
"""Load the repo as a fresh package alias for isolated tests."""
|
|
install_comfy_stubs(progress_bar_class=progress_bar_class)
|
|
unload_package(package_name)
|
|
|
|
spec = importlib.util.spec_from_file_location(
|
|
package_name,
|
|
PLUGIN_INIT_PATH,
|
|
submodule_search_locations=[str(REPO_ROOT)],
|
|
)
|
|
if spec is None or spec.loader is None:
|
|
raise RuntimeError("Failed to create an import spec for the plugin package.")
|
|
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules[package_name] = module
|
|
spec.loader.exec_module(module)
|
|
return module
|