robertvoy-ComfyUI-Distributed/tests/test_async_helpers.py

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


class _PromptQueue:
    def __init__(self):
        self.items = []

    def put(self, item):
        self.items.append(item)


def _load_async_helpers_module():
    module_path = Path(__file__).resolve().parents[1] / "utils" / "async_helpers.py"
    package_name = "dist_async_helpers_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

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

    execution_module = types.ModuleType("execution")

    async def _validate_prompt(prompt_id, prompt, partial_execution_targets):
        return (True, None, ["9"], {})

    execution_module.validate_prompt = _validate_prompt
    execution_module.SENSITIVE_EXTRA_DATA_KEYS = []
    sys.modules["execution"] = execution_module

    prompt_server = types.SimpleNamespace(
        trigger_on_prompt=lambda payload: payload,
        number=12,
        prompt_queue=_PromptQueue(),
    )
    server_module = types.ModuleType("server")
    server_module.PromptServer = types.SimpleNamespace(instance=prompt_server)
    sys.modules["server"] = server_module

    network_module = types.ModuleType(f"{package_name}.utils.network")
    network_module.get_server_loop = lambda: None
    sys.modules[f"{package_name}.utils.network"] = network_module

    spec = importlib.util.spec_from_file_location(f"{package_name}.utils.async_helpers", 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)
    return module, prompt_server


async_helpers, prompt_server = _load_async_helpers_module()


class QueuePromptPayloadTests(unittest.IsolatedAsyncioTestCase):
    async def test_insert_if_is_checked_at_insertion_after_validation_yields(self):
        module, prompt_server = _load_async_helpers_module()
        validating = asyncio.Event()
        resume = asyncio.Event()
        state = {"cancelled": False}

        async def slow_validate(prompt_id, prompt, partial_execution_targets):
            validating.set()
            await resume.wait()
            return (True, None, ["9"], {})

        module.execution.validate_prompt = slow_validate
        task = asyncio.ensure_future(module.queue_prompt_payload(
            {"9": {"class_type": "SaveImage", "inputs": {}}}, None, "client", prompt_id="run-1",
            insert_if=lambda: not state["cancelled"],
        ))
        await validating.wait()
        state["cancelled"] = True  # the queue is cleared while validation is in progress
        resume.set()
        with self.assertRaises(module.QueueInsertionCancelled):
            await task
        self.assertEqual(prompt_server.prompt_queue.items, [])
        self.assertEqual(prompt_server.number, 12)  # no queue number consumed either

        state["cancelled"] = False
        resume.set()
        prompt_id = await module.queue_prompt_payload(
            {"9": {"class_type": "SaveImage", "inputs": {}}}, None, "client", prompt_id="run-1",
            insert_if=lambda: not state["cancelled"],
        )
        self.assertEqual(prompt_id, "run-1")
        self.assertEqual([item[1] for item in prompt_server.prompt_queue.items], ["run-1"])

    async def test_queue_prompt_payload_includes_create_time_and_client_metadata(self):
        result = await async_helpers.queue_prompt_payload(
            {"1": {"class_type": "Node"}},
            workflow_meta={"id": "workflow-1"},
            client_id="client-1",
            include_queue_metadata=True,
        )

        self.assertIsInstance(result["prompt_id"], str)
        self.assertTrue(result["prompt_id"])
        self.assertEqual(result["number"], 12)
        self.assertEqual(result["node_errors"], {})

        self.assertEqual(prompt_server.number, 13)
        self.assertEqual(len(prompt_server.prompt_queue.items), 1)
        queued_item = prompt_server.prompt_queue.items[0]
        self.assertEqual(queued_item[0], 12)
        extra_data = queued_item[3]
        self.assertEqual(extra_data["client_id"], "client-1")
        self.assertIn("create_time", extra_data)
        self.assertIsInstance(extra_data["create_time"], int)
        self.assertGreater(extra_data["create_time"], 0)
        self.assertEqual(extra_data["extra_pnginfo"]["workflow"], {"id": "workflow-1"})


class ExplicitTargetTests(unittest.IsolatedAsyncioTestCase):
    def setUp(self):
        self.module, self.server = _load_async_helpers_module()
        self.prompt = {"9": {"class_type": "Output"}, "10": {"class_type": "Output"}}
        self.validate = AsyncMock(return_value=(True, None, ["9"], {}))
        self.module.execution.validate_prompt = self.validate

    async def test_targets_survive_callback_omission_and_are_queued_exactly(self):
        seen = []
        def callback(payload):
            seen.append(payload["partial_execution_targets"][:])
            payload["partial_execution_targets"].append("10")
            return {"prompt": payload["prompt"]}
        self.server.trigger_on_prompt = callback
        targets = ["9"]
        await self.module.queue_prompt_payload(self.prompt, partial_execution_targets=targets)
        self.assertEqual(seen, [["9"]])
        self.assertEqual(targets, ["9"])
        self.assertEqual(self.validate.call_args.args[2], ["9"])
        self.assertEqual(self.server.prompt_queue.items[0][4], ["9"])

    async def test_invalid_requests_never_reach_validation_or_queue(self):
        for targets in ([], "9", [9], [True], ["9", "9"], ["missing"], {}):
            with self.subTest(targets=targets):
                with self.assertRaises(self.module.PromptValidationError):
                    await self.module.queue_prompt_payload(self.prompt, partial_execution_targets=targets)
        self.validate.assert_not_awaited()
        self.assertEqual(self.server.prompt_queue.items, [])

    async def test_removed_target_after_callback_is_rejected(self):
        self.server.trigger_on_prompt = lambda payload: {"prompt": {"10": self.prompt["10"]}}
        with self.assertRaises(self.module.PromptValidationError):
            await self.module.queue_prompt_payload(self.prompt, partial_execution_targets=["9"])
        self.validate.assert_not_awaited()
        self.assertEqual(self.server.prompt_queue.items, [])

    async def test_partial_subset_extra_roots_and_target_errors_fail_closed(self):
        for roots, errors in ((["9"], {}), (["9", "10", "extra"], {}),
                              (["9", "10"], {"9": {"errors": [{"message": "bad"}]}}),
                              (["9", "10"], {"dep": {"dependent_outputs": ["10"], "errors": []}})):
            with self.subTest(roots=roots, errors=errors):
                self.validate.return_value = (True, None, roots, errors)
                with self.assertRaises(self.module.PromptValidationError):
                    await self.module.queue_prompt_payload(self.prompt, partial_execution_targets=["9", "10"])
        self.assertEqual(self.server.prompt_queue.items, [])

    async def test_none_preserves_legacy_payload_and_validation(self):
        seen = []
        self.server.trigger_on_prompt = lambda payload: seen.append(payload.copy()) or payload
        await self.module.queue_prompt_payload(self.prompt)
        self.assertNotIn("partial_execution_targets", seen[0])
        self.assertIsNone(self.validate.call_args.args[2])


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