robertvoy-ComfyUI-Distributed/tests/test_image_list_master_selection.py

"""Image-list participation uses the existing master orchestrator setting."""
import json
import sys
import unittest
from unittest.mock import AsyncMock, patch

from test_serverless_integration import PKG, load_modules


class MasterSelectionTests(unittest.IsolatedAsyncioTestCase):
    async def asyncSetUp(self):
        self.q, modules, self.ps = load_modules()
        self.modules_patch = patch.dict(sys.modules, modules)
        self.modules_patch.start()
        self.addCleanup(self.modules_patch.stop)
        self.worker = {"id": "worker", "name": "Worker", "type": "local", "enabled": True}
        self.config = {"workers": [self.worker], "settings": {}}
        self.q.load_config = lambda: self.config
        # Model the legacy helper's fallback: image-list must not inherit it.
        self.q.select_active_workers = AsyncMock(side_effect=lambda workers, ws, delegate, **kw:
                                                (workers, delegate if workers else False))
        self.preflight = AsyncMock(return_value={"worker": True})
        sys.modules[PKG + ".api.orchestration.dispatch"].preflight_image_list_workers = self.preflight
        self.q.queue_prompt_payload = AsyncMock(return_value={"prompt_id": "master-prompt", "number": 1})
        self.prompt = {
            "1": {"class_type": "EmptyImage", "inputs": {}},
            "2": {"class_type": "DistributedImageListStart", "inputs": {"images": ["1", 0]}},
            "3": {"class_type": "DistributedImageListEnd", "inputs": {"images": ["2", 0], "context": ["2", 3]}},
        }

    async def queue(self, **kwargs):
        return await self.q.orchestrate_distributed_execution(self.prompt, None, "client", **kwargs)

    def prepared_inputs(self):
        return self.q.queue_prompt_payload.call_args.args[0]["2"]["inputs"]

    async def test_default_includes_master_and_direct_workers(self):
        result = await self.queue()
        self.assertIs(self.prepared_inputs()["image_list_master_participates"], True)
        self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), ["worker"])
        self.assertEqual(result[2], 1)  # Public worker_count still counts remote participants only.
        self.preflight.assert_awaited_once_with(
            [{"id": "worker", "name": "Worker", "type": "local", "host": None, "port": 8188}],
            require_idle=False)
        self.q.dispatch_worker_prompt.assert_not_awaited()
        self.assertEqual(self.ps.distributed_pending_jobs, {})

    async def test_orchestrator_setting_excludes_master_processing(self):
        self.config["settings"]["master_delegate_only"] = True
        await self.queue()
        self.assertIs(self.prepared_inputs()["image_list_master_participates"], False)
        self.assertTrue(self.preflight.call_args.kwargs["require_idle"])
        # Upstream source stays on master even in orchestrator-only mode.
        self.assertIn("1", self.q.queue_prompt_payload.call_args.args[0])

    async def test_explicit_override_wins_over_config_in_both_directions(self):
        for configured, requested in [(True, False), (False, True)]:
            with self.subTest(configured=configured, requested=requested):
                self.config["settings"]["master_delegate_only"] = configured
                await self.queue(delegate_master=requested)
                self.assertIs(self.prepared_inputs()["image_list_master_participates"], not requested)

    async def test_master_only_with_explicit_empty_worker_selection(self):
        result = await self.queue(enabled_worker_ids=[])
        self.assertIs(self.prepared_inputs()["image_list_master_participates"], True)
        self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), [])
        self.assertEqual(result[2], 0)
        self.preflight.assert_not_awaited()

    async def test_master_only_when_no_workers_are_configured(self):
        self.config["workers"] = []
        await self.queue()
        self.assertIs(self.prepared_inputs()["image_list_master_participates"], True)

    async def test_unavailable_workers_do_not_prevent_enabled_master(self):
        self.q.select_active_workers.return_value = ([], False)
        self.q.select_active_workers.side_effect = None
        self.preflight.return_value = {}
        await self.queue()
        self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), [])
        self.assertIs(self.prepared_inputs()["image_list_master_participates"], True)

    async def test_busy_workers_do_not_prevent_enabled_master(self):
        self.preflight.return_value = {}
        await self.queue()
        self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), [])

    async def test_orchestrator_without_workers_fails_before_queue(self):
        self.config["settings"]["master_delegate_only"] = True
        with self.assertRaisesRegex(RuntimeError, "orchestrator|Orchestrator"):
            await self.queue(enabled_worker_ids=[])
        self.q.queue_prompt_payload.assert_not_awaited()

    async def test_orchestrator_never_inherits_legacy_offline_master_fallback(self):
        self.config["settings"]["master_delegate_only"] = True
        self.q.select_active_workers.side_effect = None
        self.q.select_active_workers.return_value = ([], False)
        self.preflight.return_value = {}
        with self.assertRaisesRegex(RuntimeError, "orchestrator|Orchestrator"):
            await self.queue()
        self.q.queue_prompt_payload.assert_not_awaited()

    async def test_serverless_only_uses_master_without_provider_dispatch(self):
        self.config["workers"] = [{"id": "pool", "type": "serverless", "enabled": True}]
        result = await self.queue()
        self.assertIs(self.prepared_inputs()["image_list_master_participates"], True)
        self.assertEqual(result[2], 0)
        self.q.dispatch_worker_prompt.assert_not_awaited()

    async def test_incompatible_worker_is_not_silently_hidden_by_master(self):
        self.preflight.side_effect = RuntimeError("upgrade worker")
        with self.assertRaisesRegex(RuntimeError, "upgrade worker"):
            await self.queue()
        self.q.queue_prompt_payload.assert_not_awaited()