JasonHoku-ComfyUI-Ultimate-Auto-Sampler-Config-Grid-Testing-Suite/conftest.py

"""
Root-level pytest configuration for standalone testing outside ComfyUI.

Installs ComfyUI stubs and a custom import hook so that the package's
__init__.py (which pytest's Package.setup() always imports) loads cleanly
without a live ComfyUI runtime.
"""
import sys
import types
import os
import importlib
import importlib.abc
import importlib.machinery


# Run immediately at conftest import time (before pytest_configure)
_NODE_ROOT = os.path.dirname(os.path.abspath(__file__))
_PKG_NAME = os.path.basename(_NODE_ROOT)  # "ComfyUI-Ultimate-Auto-Sampler-Config-Grid-Testing-Suite"


def _install_stubs():
    """Install all stubs needed to import the package __init__.py cleanly."""

    # 1. External deps
    for name in [
        "torch", "torchvision", "torchaudio",
        "folder_paths",
        "comfy", "comfy.utils", "comfy.sd", "comfy.model_management",
        "nodes", "aiohttp",
    ]:
        if name not in sys.modules:
            sys.modules[name] = types.ModuleType(name)

    # aiohttp.web
    if "aiohttp.web" not in sys.modules:
        web = types.ModuleType("aiohttp.web")
        web.Response = type("Response", (), {"__init__": lambda s, *a, **kw: None})
        web.json_response = lambda *a, **kw: None
        web.RouteTableDef = type("RouteTableDef", (), {
            "get": lambda s, *a, **kw: (lambda f: f),
            "post": lambda s, *a, **kw: (lambda f: f),
            "delete": lambda s, *a, **kw: (lambda f: f),
            "put": lambda s, *a, **kw: (lambda f: f),
        })
        sys.modules["aiohttp.web"] = web
        sys.modules["aiohttp"].web = web  # type: ignore[attr-defined]

    # folder_paths.get_output_directory
    fp = sys.modules["folder_paths"]
    if not hasattr(fp, "get_output_directory"):
        fp.get_output_directory = lambda: os.path.join(_NODE_ROOT, "_test_output")  # type: ignore

    # server.PromptServer
    if "server" not in sys.modules:
        sys.modules["server"] = types.ModuleType("server")
    srv = sys.modules["server"]
    if not hasattr(srv, "PromptServer"):
        _noop_deco = staticmethod(lambda *a, **kw: (lambda f: f))
        _routes = type("_Routes", (), {
            "get": _noop_deco, "post": _noop_deco,
            "delete": _noop_deco, "put": _noop_deco,
        })()
        _ps_instance = type("_PS", (), {"routes": _routes})()
        srv.PromptServer = type("PromptServer", (), {"instance": _ps_instance})  # type: ignore

    # 2. Sub-module stubs (for relative imports in __init__.py)
    _subs = {
        "sampler_node": {"SamplerGridTester": type("SamplerGridTester", (), {})},
        "dashboard_node": {"SamplerConfigDashboardViewer": type("SamplerConfigDashboardViewer", (), {})},
        "html_generator": {"get_html_template": lambda *a, **kw: ""},
        "config_builder_node": {"UltimateConfigBuilder": type("UltimateConfigBuilder", (), {})},
        "json_text_node": {"SmartJSONTextNode": type("SmartJSONTextNode", (), {})},
        "metadata_packer": {"pack_metadata_into_image": lambda *a, **kw: None},
        "directory_scanner": {
            "scan_directory_for_images": lambda *a, **kw: ([], {}),
            "MEDIA_EXTENSIONS": set(),
        },
        "distribution_routes": {},
        "upscale_runner": {
            "start_upscale_job": lambda *a, **kw: (None, None),
            "get_upscale_status": lambda *a, **kw: {},
            "cancel_upscale_job": lambda *a, **kw: False,
        },
    }

    for bare, attrs in _subs.items():
        fq = f"{_PKG_NAME}.{bare}"
        if bare not in sys.modules:
            stub = types.ModuleType(bare)
            sys.modules[bare] = stub
        if fq not in sys.modules:
            sys.modules[fq] = sys.modules[bare]
        for k, v in attrs.items():
            setattr(sys.modules[bare], k, v)
            setattr(sys.modules[fq], k, v)

    # 3. Pre-register the package itself
    if _PKG_NAME not in sys.modules:
        pkg_stub = types.ModuleType(_PKG_NAME)
        pkg_stub.__path__ = [_NODE_ROOT]
        pkg_stub.__package__ = _PKG_NAME
        sys.modules[_PKG_NAME] = pkg_stub

    # 4. Ensure node root is on sys.path
    if _NODE_ROOT not in sys.path:
        sys.path.insert(0, _NODE_ROOT)


class _PackageInitLoader(importlib.abc.Loader):
    """
    Custom loader for the package __init__.py.
    Sets __package__ correctly before executing the module so that relative
    imports (from .sampler_node import X) resolve via sys.modules.
    """

    def __init__(self, path):
        self._path = path

    def create_module(self, spec):
        return None  # use default semantics

    def exec_module(self, module):
        module.__file__ = self._path
        module.__package__ = _PKG_NAME
        module.__name__ = _PKG_NAME
        module.__path__ = [_NODE_ROOT]  # makes relative imports resolve parent pkg
        # Point sys.modules to this module object so relative imports find it
        sys.modules[_PKG_NAME] = module
        with open(self._path, "rb") as fh:
            src = fh.read()
        code = compile(src, self._path, "exec")
        exec(code, module.__dict__)  # noqa: S102  – intentional dynamic exec for stub loading


class _PackageInitFinder(importlib.abc.MetaPathFinder):
    """
    Intercepts any attempt to import the package __init__.py and routes it
    through _PackageInitLoader so relative imports get the correct __package__.
    """
    _INIT_PATH = os.path.join(_NODE_ROOT, "__init__.py")

    def find_spec(self, fullname, path, target=None):
        # In importlib mode, pytest imports __init__.py as module name "__init__"
        # (module_name_from_path returns "__init__" for package root __init__.py).
        # We intercept that import and also the top-level package name import.
        if fullname in ("__init__", _PKG_NAME):
            if target is not None:
                # target is the path pytest wants to load; check it's our __init__.py
                try:
                    if os.path.normcase(str(target.__spec__.origin)) != os.path.normcase(self._INIT_PATH):
                        return None
                except Exception:
                    pass
            spec = importlib.machinery.ModuleSpec(
                fullname,
                _PackageInitLoader(self._INIT_PATH),
                origin=self._INIT_PATH,
            )
            spec.submodule_search_locations = [_NODE_ROOT]
            return spec
        return None


# Install stubs + hook immediately so they're ready before pytest does anything
_install_stubs()
sys.meta_path.insert(0, _PackageInitFinder())


def pytest_configure(config):
    """Re-ensure stubs are in place after config initialises."""
    _install_stubs()