robertvoy-ComfyUI-Distributed/tests/test_dispatch_selection.py

import asyncio
import importlib.util
import sys
import types
import unittest
from pathlib import Path
from unittest.mock import patch


def _load_dispatch_module():
    module_path = Path(__file__).resolve().parents[1] / "api" / "orchestration" / "dispatch.py"

    package_name = "dist_dispatch_testpkg"
    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

    orch_pkg = types.ModuleType(f"{package_name}.api.orchestration")
    orch_pkg.__path__ = []
    sys.modules[f"{package_name}.api.orchestration"] = orch_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

    network_module = types.ModuleType(f"{package_name}.utils.network")
    network_module.build_worker_url = lambda *_args, **_kwargs: "http://example.invalid"

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

    network_module.probe_worker = _probe_worker

    async def _fake_session():
        raise RuntimeError("get_client_session should be mocked in these tests")

    network_module.get_client_session = _fake_session
    sys.modules[f"{package_name}.utils.network"] = network_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 _ClientConnectorError(Exception):
            pass

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

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

        class _ClientSession:
            def __init__(self, *args, **kwargs):
                self.closed = False

            async def close(self):
                self.closed = True

        aiohttp_module.ClientTimeout = _ClientTimeout
        aiohttp_module.ClientConnectorError = _ClientConnectorError
        aiohttp_module.WSMsgType = _WSMsgType
        aiohttp_module.TCPConnector = _TCPConnector
        aiohttp_module.ClientSession = _ClientSession
        aiohttp_module.web = types.SimpleNamespace(
            json_response=lambda payload, status=200: {"payload": payload, "status": status}
        )
        sys.modules["aiohttp"] = aiohttp_module

    spec = importlib.util.spec_from_file_location(
        f"{package_name}.api.orchestration.dispatch",
        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)
    if created_aiohttp_stub:
        sys.modules.pop("aiohttp", None)
    return module


dispatch = _load_dispatch_module()


class DispatchSelectionTests(unittest.IsolatedAsyncioTestCase):
    async def test_select_active_workers_filters_offline(self):
        workers = [
            {"id": "w1", "name": "Worker 1"},
            {"id": "w2", "name": "Worker 2"},
            {"id": "w3", "name": "Worker 3"},
        ]

        async def fake_probe(worker):
            return worker["id"] != "w2"

        with patch.object(dispatch, "worker_is_active", side_effect=fake_probe):
            active_workers, delegate_master = await dispatch.select_active_workers(
                workers,
                use_websocket=False,
                delegate_master=False,
                probe_concurrency=3,
            )

        self.assertEqual([w["id"] for w in active_workers], ["w1", "w3"])
        self.assertFalse(delegate_master)

    async def test_select_active_workers_disables_delegate_when_all_offline(self):
        workers = [{"id": "w1", "name": "Worker 1"}]

        async def fake_probe(_worker):
            return False

        with patch.object(dispatch, "worker_is_active", side_effect=fake_probe):
            active_workers, delegate_master = await dispatch.select_active_workers(
                workers,
                use_websocket=False,
                delegate_master=True,
                probe_concurrency=1,
            )

        self.assertEqual(active_workers, [])
        self.assertFalse(delegate_master)

    async def test_select_active_workers_uses_websocket_probe_when_enabled(self):
        workers = [{"id": "w1", "name": "Worker 1"}, {"id": "w2", "name": "Worker 2"}]

        async def fake_http_probe(_worker):
            return False

        async def fake_ws_probe(_worker):
            return True

        with patch.object(dispatch, "worker_is_active", side_effect=fake_http_probe) as http_probe, patch.object(
            dispatch,
            "worker_ws_is_active",
            side_effect=fake_ws_probe,
        ) as ws_probe:
            active_workers, _ = await dispatch.select_active_workers(
                workers,
                use_websocket=True,
                delegate_master=False,
                probe_concurrency=4,
            )

        self.assertEqual([w["id"] for w in active_workers], ["w1", "w2"])
        self.assertEqual(ws_probe.call_count, 2)
        self.assertEqual(http_probe.call_count, 0)

    async def test_probe_concurrency_is_bounded(self):
        workers = [{"id": f"w{i}", "name": f"Worker {i}"} for i in range(6)]
        state = {"in_flight": 0, "max_in_flight": 0}

        async def fake_probe(_worker):
            state["in_flight"] += 1
            state["max_in_flight"] = max(state["max_in_flight"], state["in_flight"])
            await asyncio.sleep(0.01)
            state["in_flight"] -= 1
            return True

        with patch.object(dispatch, "worker_is_active", side_effect=fake_probe):
            active_workers, _ = await dispatch.select_active_workers(
                workers,
                use_websocket=False,
                delegate_master=False,
                probe_concurrency=2,
            )

        self.assertEqual(len(active_workers), len(workers))
        self.assertLessEqual(state["max_in_flight"], 2)
        self.assertGreaterEqual(state["max_in_flight"], 2)


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