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

import importlib.util
import os
import sys
import tempfile
import types
import unittest
from collections import deque
from pathlib import Path
from unittest.mock import patch


class _FakeResponse:
    def __init__(self, payload, status=200):
        self.payload = payload
        self.status = status


class _FakeRequest:
    def __init__(self, payload=None, match_info=None, query=None):
        self._payload = payload
        self.match_info = match_info or {}
        self.query = query or {}

    async def json(self):
        return self._payload


class _FakeHTTPClientResponse:
    def __init__(self, payload, status=200):
        self._payload = payload
        self.status = status

    async def __aenter__(self):
        return self

    async def __aexit__(self, _exc_type, _exc, _tb):
        return False

    async def json(self):
        return self._payload

    async def text(self):
        return str(self._payload)


class _FakeHTTPClientSession:
    def __init__(self, payload, status=200):
        self._payload = payload
        self._status = status
        self.calls = []

    def get(self, url, params=None, timeout=None):
        self.calls.append({"url": url, "params": params, "timeout": timeout})
        return _FakeHTTPClientResponse(self._payload, status=self._status)


class _DummyWorkerManager:
    def __init__(self):
        self.processes = {}

    def launch_worker(self, worker):
        worker_id = str(worker["id"])
        self.processes[worker_id] = {
            "pid": 12345,
            "log_file": f"/tmp/distributed_worker_{worker_id}.log",
            "process": None,
        }
        return 12345

    def _is_process_running(self, _pid):
        return False

    def save_processes(self):
        return None

    def stop_worker(self, _worker_id):
        return True, "Stopped"

    def get_managed_workers(self):
        return []


class _ImmediateLoop:
    async def run_in_executor(self, _executor, func, *args):
        return func(*args)


def _load_worker_routes_module():
    module_path = Path(__file__).resolve().parents[2] / "api" / "worker_routes.py"
    package_name = "dist_api_worker_testpkg"

    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

    api_pkg = types.ModuleType(f"{package_name}.api")
    api_pkg.__path__ = []
    sys.modules[f"{package_name}.api"] = api_pkg

    utils_pkg = types.ModuleType(f"{package_name}.utils")
    utils_pkg.__path__ = []
    sys.modules[f"{package_name}.utils"] = utils_pkg

    workers_pkg = types.ModuleType(f"{package_name}.workers")
    workers_pkg.__path__ = [str(module_path.parents[1] / "workers")]
    workers_pkg.get_worker_manager = lambda: _DummyWorkerManager()
    sys.modules[f"{package_name}.workers"] = workers_pkg

    detection_module = types.ModuleType(f"{package_name}.workers.detection")
    detection_module.is_local_worker = lambda *_args, **_kwargs: True
    detection_module.is_same_physical_host = lambda *_args, **_kwargs: True
    detection_module.get_machine_id = lambda: "machine-id"
    detection_module.is_docker_environment = lambda: False
    detection_module.is_runpod_environment = lambda: False
    detection_module.get_comms_channel = lambda *_args, **_kwargs: "lan"
    sys.modules[f"{package_name}.workers.detection"] = detection_module

    created_aiohttp_stub = False
    if "aiohttp" not in sys.modules:
        created_aiohttp_stub = True
        aiohttp_module = types.ModuleType("aiohttp")

        class _ClientTimeout:
            def __init__(self, total=None):
                self.total = total

        class _WSMsgType:
            TEXT = "TEXT"
            ERROR = "ERROR"
            CLOSED = "CLOSED"

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

            async def prepare(self, _request):
                return None

            async def send_json(self, _payload):
                return None

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

        aiohttp_module.ClientTimeout = _ClientTimeout
        aiohttp_module.WSMsgType = _WSMsgType
        aiohttp_module.web = types.SimpleNamespace(
            json_response=lambda payload, status=200: _FakeResponse(payload, status=status),
            WebSocketResponse=_WebSocketResponse,
        )
        sys.modules["aiohttp"] = aiohttp_module

    class _Routes:
        def get(self, _path):
            def _decorator(fn):
                return fn
            return _decorator

        def post(self, _path):
            def _decorator(fn):
                return fn
            return _decorator

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

    created_torch_stub = False
    if "torch" not in sys.modules:
        created_torch_stub = True
        torch_module = types.ModuleType("torch")
        torch_module.cuda = types.SimpleNamespace(
            is_available=lambda: False,
            empty_cache=lambda: None,
            ipc_collect=lambda: None,
            current_device=lambda: 0,
            device_count=lambda: 0,
        )
        sys.modules["torch"] = torch_module

    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

    config_module = types.ModuleType(f"{package_name}.utils.config")
    config_module.load_config = lambda: {"workers": []}
    sys.modules[f"{package_name}.utils.config"] = config_module

    network_module = types.ModuleType(f"{package_name}.utils.network")

    async def _handle_api_error(_request, error, status=500):
        return _FakeResponse({"status": "error", "message": str(error)}, status=status)

    network_module.handle_api_error = _handle_api_error
    network_module.normalize_host = lambda value: value
    network_module.build_worker_url = lambda worker, endpoint="": f"http://localhost:{worker.get('port', 8188)}{endpoint}"

    async def _probe_worker(*_args, **_kwargs):
        return None

    network_module.probe_worker = _probe_worker

    async def _get_client_session():
        raise RuntimeError("not used in these tests")

    network_module.get_client_session = _get_client_session
    network_module.get_server_port = lambda: 8189
    sys.modules[f"{package_name}.utils.network"] = network_module

    constants_module = types.ModuleType(f"{package_name}.utils.constants")
    constants_module.CHUNK_SIZE = 8192
    sys.modules[f"{package_name}.utils.constants"] = constants_module

    async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")

    async def _queue_prompt_payload(*_args, **_kwargs):
        return "prompt-id"

    async_helpers_module.queue_prompt_payload = _queue_prompt_payload

    class _PromptValidationError(RuntimeError):
        def __init__(self, message="invalid prompt", validation_error=None, node_errors=None):
            super().__init__(message)
            self.validation_error = validation_error if isinstance(validation_error, dict) else {}
            self.node_errors = node_errors if isinstance(node_errors, dict) else {}

    async_helpers_module.PromptValidationError = _PromptValidationError
    sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module

    schemas_module = types.ModuleType(f"{package_name}.api.schemas")

    def _require_fields(data, *fields):
        missing = []
        for field in fields:
            value = data.get(field) if isinstance(data, dict) else None
            if value is None or (isinstance(value, str) and not value.strip()):
                missing.append(field)
        return missing

    def _validate_worker_id(worker_id, config):
        return any(str(worker.get("id")) == str(worker_id) for worker in config.get("workers", []))

    schemas_module.require_fields = _require_fields
    schemas_module.validate_worker_id = _validate_worker_id
    sys.modules[f"{package_name}.api.schemas"] = schemas_module

    spec = importlib.util.spec_from_file_location(f"{package_name}.api.worker_routes", 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)
    # Keep this module's response double independent of aiohttp import order.
    module.web = types.SimpleNamespace(json_response=lambda payload, status=200: _FakeResponse(payload, status))

    if created_aiohttp_stub:
        sys.modules.pop("aiohttp", None)
    if created_torch_stub:
        sys.modules.pop("torch", None)

    return module


worker_routes = _load_worker_routes_module()


class WorkerRoutesTests(unittest.IsolatedAsyncioTestCase):
    async def test_network_info_reports_actual_master_and_available_ports(self):
        with patch.object(worker_routes, "_get_cuda_info", return_value=(0, 1, 4)), \
             patch.object(worker_routes, "get_network_ips", return_value=["127.0.0.1"]), \
             patch.object(worker_routes, "load_config", return_value={"workers": []}), \
             patch.object(worker_routes, "get_server_port", return_value=8189), \
             patch.object(worker_routes, "allocate_worker_ports", return_value=[8191, 8192, 8193]) as allocate:
            response = await worker_routes.get_network_info_endpoint(_FakeRequest())
        self.assertEqual(response.payload["master_port"], 8189)
        self.assertEqual(response.payload["local_worker_ports"], [8191, 8192, 8193])
        self.assertEqual(response.payload["cuda_device_count"], 4)
        allocate.assert_called_once_with(8189, [], 3)

    async def test_network_info_does_not_reallocate_existing_configuration(self):
        with patch.object(worker_routes, "_get_cuda_info", return_value=(0, 4, 4)), \
             patch.object(worker_routes, "get_network_ips", return_value=[]), \
             patch.object(worker_routes, "load_config", return_value={"workers": [{"id": "existing", "port": 8189}]}), \
             patch.object(worker_routes, "allocate_worker_ports") as allocate:
            response = await worker_routes.get_network_info_endpoint(_FakeRequest())
        self.assertEqual(response.payload["local_worker_ports"], [])
        allocate.assert_not_called()

    async def test_launch_conflict_is_a_clear_client_error(self):
        manager = _DummyWorkerManager()
        config = {"workers": [{"id": "worker-a", "name": "Worker A", "port": 8189}]}
        with patch.object(worker_routes, "get_worker_manager", return_value=manager), \
             patch.object(worker_routes, "load_config", return_value=config), \
             patch.object(manager, "launch_worker", side_effect=ValueError("Worker port 8189 conflicts with the master port 8189")):
            response = await worker_routes.launch_worker_endpoint(_FakeRequest({"worker_id": "worker-a"}))
        self.assertEqual(response.status, 400)
        self.assertIn("master port 8189", response.payload["message"])
        self.assertEqual(manager.processes, {})

    async def test_launch_worker_valid_id_returns_200(self):
        manager = _DummyWorkerManager()
        config = {"workers": [{"id": "worker-a", "name": "Worker A", "port": 8188}]}
        request = _FakeRequest({"worker_id": "worker-a"})

        with patch.object(worker_routes, "get_worker_manager", return_value=manager), patch.object(
            worker_routes, "load_config", return_value=config
        ), patch.object(
            worker_routes.asyncio, "get_running_loop", return_value=_ImmediateLoop()
        ):
            response = await worker_routes.launch_worker_endpoint(request)

        self.assertEqual(response.status, 200)
        self.assertEqual(response.payload.get("status"), "success")
        self.assertEqual(response.payload.get("pid"), 12345)

    async def test_launch_worker_unknown_id_returns_404(self):
        manager = _DummyWorkerManager()
        config = {"workers": [{"id": "worker-a", "name": "Worker A", "port": 8188}]}
        request = _FakeRequest({"worker_id": "missing-worker"})

        with patch.object(worker_routes, "get_worker_manager", return_value=manager), patch.object(
            worker_routes, "load_config", return_value=config
        ):
            response = await worker_routes.launch_worker_endpoint(request)

        self.assertEqual(response.status, 404)
        self.assertIn("not found", response.payload.get("message", "").lower())

    async def test_worker_log_returns_content_json(self):
        manager = _DummyWorkerManager()
        with tempfile.NamedTemporaryFile("w", delete=False, encoding="utf-8") as handle:
            handle.write("line-1\nline-2\nline-3\n")
            log_path = handle.name

        manager.processes["worker-a"] = {
            "pid": 9999,
            "log_file": log_path,
            "process": None,
        }

        request = _FakeRequest(match_info={"worker_id": "worker-a"}, query={"lines": "2"})
        try:
            with patch.object(worker_routes, "get_worker_manager", return_value=manager), patch.object(
                worker_routes.asyncio, "get_running_loop", return_value=_ImmediateLoop()
            ):
                response = await worker_routes.get_worker_log_endpoint(request)
        finally:
            if os.path.exists(log_path):
                os.remove(log_path)

        self.assertEqual(response.status, 200)
        self.assertEqual(response.payload.get("status"), "success")
        self.assertIn("content", response.payload)
        self.assertIn("line-3", response.payload["content"])

    async def test_local_log_reads_memory_buffer(self):
        request = _FakeRequest(query={"lines": "2"})
        fake_logs = deque(
            [
                {"m": "line-1\n"},
                {"m": "line-2\n"},
                {"m": "line-3\n"},
            ],
            maxlen=300,
        )
        app_module = types.ModuleType("app")
        app_module.__path__ = []
        logger_module = types.ModuleType("app.logger")
        logger_module.get_logs = lambda: fake_logs
        app_module.logger = logger_module

        with patch.dict(sys.modules, {"app": app_module, "app.logger": logger_module}):
            response = await worker_routes.get_local_log_endpoint(request)

        self.assertEqual(response.status, 200)
        self.assertEqual(response.payload.get("status"), "success")
        self.assertEqual(response.payload.get("source"), "memory")
        self.assertEqual(response.payload.get("entries"), 2)
        self.assertIn("line-3", response.payload.get("content", ""))

    async def test_remote_worker_log_proxies_to_worker_local_log_endpoint(self):
        config = {
            "workers": [
                {
                    "id": "worker-remote",
                    "name": "Remote Worker",
                    "host": "worker.example.com",
                    "port": 8188,
                    "type": "remote",
                }
            ]
        }
        request = _FakeRequest(match_info={"worker_id": "worker-remote"}, query={"lines": "120"})
        proxied_payload = {
            "status": "success",
            "content": "remote-log-content\n",
            "entries": 1,
            "source": "memory",
            "truncated": False,
            "lines_shown": 1,
        }
        fake_session = _FakeHTTPClientSession(proxied_payload)

        async def _fake_get_client_session():
            return fake_session

        with patch.object(worker_routes, "load_config", return_value=config), patch.object(
            worker_routes, "get_client_session", side_effect=_fake_get_client_session
        ):
            response = await worker_routes.get_remote_worker_log_endpoint(request)

        self.assertEqual(response.status, 200)
        self.assertEqual(response.payload.get("content"), "remote-log-content\n")
        self.assertEqual(len(fake_session.calls), 1)
        self.assertEqual(fake_session.calls[0]["params"], {"lines": "120"})
        self.assertTrue(fake_session.calls[0]["url"].endswith("/distributed/local_log"))

    async def test_remote_worker_log_rejects_local_workers(self):
        config = {"workers": [{"id": "worker-local", "name": "Local Worker", "port": 8188}]}
        request = _FakeRequest(match_info={"worker_id": "worker-local"})

        with patch.object(worker_routes, "load_config", return_value=config):
            response = await worker_routes.get_remote_worker_log_endpoint(request)

        self.assertEqual(response.status, 400)
        self.assertIn("local", response.payload.get("message", "").lower())


if __name__ == "__main__":
    unittest.main()