robertvoy-ComfyUI-Distributed/tests/api/harness.py

"""Shared stubs/utilities for API route module unit tests."""
from __future__ import annotations

import sys
import types
from typing import Any, Callable


class RoutesStub:
    """Simple decorator-style route registrar used in unit tests."""

    def get(self, _path: str):
        return lambda fn: fn

    def post(self, _path: str):
        return lambda fn: fn


class ClientTimeoutStub:
    def __init__(self, total: float | None = None):
        self.total = total


class WSMsgTypeStub:
    TEXT = "TEXT"
    ERROR = "ERROR"
    CLOSED = "CLOSED"


class WebSocketResponseStub:
    def __init__(self, *args: Any, **kwargs: Any):
        self.args = args
        self.kwargs = kwargs

    async def prepare(self, _request: Any) -> None:
        return None

    async def send_json(self, _payload: Any) -> None:
        return None

    def __aiter__(self):
        async def _empty():
            if False:
                yield None

        return _empty()


class FormDataStub:
    def add_field(self, *_args: Any, **_kwargs: Any) -> None:
        return None


def reset_package_namespace(package_name: str) -> None:
    """Clear test package modules from sys.modules."""
    prefix = f"{package_name}."
    for mod_name in list(sys.modules):
        if mod_name == package_name or mod_name.startswith(prefix):
            del sys.modules[mod_name]


def ensure_namespace_package(module_name: str) -> types.ModuleType:
    """Create a namespace-like package module in sys.modules."""
    module = types.ModuleType(module_name)
    module.__path__ = []
    sys.modules[module_name] = module
    return module


def bootstrap_test_package(
    package_name: str,
    *,
    with_api: bool = True,
    with_utils: bool = True,
    with_workers: bool = False,
    with_upscale: bool = False,
    with_orchestration: bool = False,
) -> None:
    """Create a clean package scaffold for module-level import tests."""
    reset_package_namespace(package_name)
    ensure_namespace_package(package_name)

    if with_api:
        ensure_namespace_package(f"{package_name}.api")
    if with_utils:
        ensure_namespace_package(f"{package_name}.utils")
    if with_workers:
        ensure_namespace_package(f"{package_name}.workers")
    if with_upscale:
        ensure_namespace_package(f"{package_name}.upscale")
    if with_orchestration:
        ensure_namespace_package(f"{package_name}.api.orchestration")


def install_server_stub(prompt_server_instance: Any | None = None) -> Any:
    """Install a minimal `server` module exposing PromptServer.instance.routes."""
    if prompt_server_instance is None:
        prompt_server_instance = types.SimpleNamespace()
    if not hasattr(prompt_server_instance, "routes"):
        prompt_server_instance.routes = RoutesStub()

    server_module = types.ModuleType("server")
    server_module.PromptServer = types.SimpleNamespace(instance=prompt_server_instance)
    sys.modules["server"] = server_module
    return prompt_server_instance


def install_aiohttp_stub(
    json_response_factory: Callable[..., Any],
) -> bool:
    """Install/augment `aiohttp` with common members used by API tests."""
    created = False
    if "aiohttp" not in sys.modules:
        created = True
        aiohttp_module = types.ModuleType("aiohttp")
        sys.modules["aiohttp"] = aiohttp_module

    aiohttp_module = sys.modules["aiohttp"]
    if not hasattr(aiohttp_module, "ClientTimeout"):
        aiohttp_module.ClientTimeout = ClientTimeoutStub
    if not hasattr(aiohttp_module, "WSMsgType"):
        aiohttp_module.WSMsgType = WSMsgTypeStub
    if not hasattr(aiohttp_module, "FormData"):
        aiohttp_module.FormData = FormDataStub

    web_obj = getattr(aiohttp_module, "web", None)
    if web_obj is None:
        web_obj = types.SimpleNamespace()
        aiohttp_module.web = web_obj
    web_obj.json_response = json_response_factory
    if not hasattr(web_obj, "WebSocketResponse"):
        web_obj.WebSocketResponse = WebSocketResponseStub

    return created


def cleanup_optional_module(module_name: str, created: bool) -> None:
    """Remove stubbed module only when this test created it."""
    if created:
        sys.modules.pop(module_name, None)


def install_request_guards_stub(package_name: str) -> None:
    """Install permissive request guard for isolated route tests."""
    request_guards_module = types.ModuleType(f"{package_name}.api.request_guards")

    async def _authorization_error_or_none(_request):
        return None

    request_guards_module.authorization_error_or_none = _authorization_error_or_none
    sys.modules[f"{package_name}.api.request_guards"] = request_guards_module


def install_endpoint_policy_passthrough_stub(package_name: str) -> None:
    """Install endpoint policy stub that executes operation directly."""
    endpoint_policy_module = types.ModuleType(f"{package_name}.api.endpoint_policy")

    async def _run_authorized_endpoint(_request, operation, unexpected_status=500):
        _ = unexpected_status
        return await operation()

    endpoint_policy_module.run_authorized_endpoint = _run_authorized_endpoint
    sys.modules[f"{package_name}.api.endpoint_policy"] = endpoint_policy_module