230 lines
8.2 KiB
Python
230 lines
8.2 KiB
Python
import importlib.util
|
|
import sys
|
|
import types
|
|
import unittest
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
|
|
class _FakeResponse:
|
|
def __init__(self, payload, status=200):
|
|
self.payload = payload
|
|
self.status = status
|
|
|
|
|
|
class _FakeRequest:
|
|
def __init__(self, payload):
|
|
self._payload = payload
|
|
|
|
async def json(self):
|
|
return self._payload
|
|
|
|
|
|
def _load_job_routes_module():
|
|
module_path = Path(__file__).resolve().parents[2] / "api" / "job_routes.py"
|
|
package_name = "dist_api_queue_testpkg"
|
|
|
|
# Reset package namespace to avoid stale module state across test runs.
|
|
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
|
|
|
|
# aiohttp.web stub
|
|
created_aiohttp_stub = False
|
|
if "aiohttp" not in sys.modules:
|
|
created_aiohttp_stub = True
|
|
aiohttp_module = types.ModuleType("aiohttp")
|
|
aiohttp_module.web = types.SimpleNamespace(
|
|
json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
|
|
)
|
|
sys.modules["aiohttp"] = aiohttp_module
|
|
|
|
# server module stub with route decorators
|
|
class _Routes:
|
|
def get(self, _path):
|
|
def _decorator(fn):
|
|
return fn
|
|
return _decorator
|
|
|
|
def post(self, _path):
|
|
def _decorator(fn):
|
|
return fn
|
|
return _decorator
|
|
|
|
prompt_server_instance = types.SimpleNamespace(
|
|
routes=_Routes(),
|
|
distributed_jobs_lock=None,
|
|
distributed_pending_jobs={},
|
|
)
|
|
server_module = types.ModuleType("server")
|
|
server_module.PromptServer = types.SimpleNamespace(instance=prompt_server_instance)
|
|
sys.modules["server"] = server_module
|
|
|
|
# torch stub (only needed to satisfy import)
|
|
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,
|
|
)
|
|
sys.modules["torch"] = torch_module
|
|
|
|
# PIL stub (only needed to satisfy import)
|
|
created_pil_stub = False
|
|
if "PIL" not in sys.modules:
|
|
created_pil_stub = True
|
|
pil_module = types.ModuleType("PIL")
|
|
image_module = types.ModuleType("PIL.Image")
|
|
pil_module.Image = image_module
|
|
sys.modules["PIL"] = pil_module
|
|
sys.modules["PIL.Image"] = image_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
|
|
|
|
image_module = types.ModuleType(f"{package_name}.utils.image")
|
|
image_module.pil_to_tensor = lambda *_args, **_kwargs: None
|
|
image_module.ensure_contiguous = lambda tensor: tensor
|
|
sys.modules[f"{package_name}.utils.image"] = image_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
|
|
sys.modules[f"{package_name}.utils.network"] = network_module
|
|
|
|
constants_module = types.ModuleType(f"{package_name}.utils.constants")
|
|
constants_module.MEMORY_CLEAR_DELAY = 0.0
|
|
constants_module.JOB_INIT_GRACE_PERIOD = 10.0
|
|
sys.modules[f"{package_name}.utils.constants"] = constants_module
|
|
|
|
async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
|
|
async_helpers_module.queue_prompt_payload = AsyncMock(return_value="prompt_local")
|
|
sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
|
|
|
|
queue_orchestration_module = types.ModuleType(f"{package_name}.api.queue_orchestration")
|
|
queue_orchestration_module.orchestrate_distributed_execution = AsyncMock(return_value=("prompt_dist", 1))
|
|
sys.modules[f"{package_name}.api.queue_orchestration"] = queue_orchestration_module
|
|
|
|
@dataclass(frozen=True)
|
|
class _QueuePayload:
|
|
prompt: dict
|
|
workflow_meta: object
|
|
client_id: str
|
|
delegate_master: object
|
|
enabled_worker_ids: list
|
|
auto_prepare: bool
|
|
trace_execution_id: object
|
|
|
|
def _parse_queue_request_payload(data):
|
|
if not isinstance(data, dict):
|
|
raise ValueError("Expected a JSON object body")
|
|
prompt = data.get("prompt")
|
|
if not isinstance(prompt, dict):
|
|
raise ValueError("Field 'prompt' must be an object")
|
|
enabled = data.get("enabled_worker_ids")
|
|
if not isinstance(enabled, list):
|
|
raise ValueError("enabled_worker_ids required")
|
|
client_id = data.get("client_id")
|
|
if not isinstance(client_id, str) or not client_id.strip():
|
|
raise ValueError("client_id required")
|
|
return _QueuePayload(
|
|
prompt=prompt,
|
|
workflow_meta=data.get("workflow"),
|
|
client_id=client_id,
|
|
delegate_master=data.get("delegate_master"),
|
|
enabled_worker_ids=enabled,
|
|
auto_prepare=bool(data.get("auto_prepare", True)),
|
|
trace_execution_id=data.get("trace_execution_id"),
|
|
)
|
|
|
|
queue_request_module = types.ModuleType(f"{package_name}.api.queue_request")
|
|
queue_request_module.parse_queue_request_payload = _parse_queue_request_payload
|
|
sys.modules[f"{package_name}.api.queue_request"] = queue_request_module
|
|
|
|
spec = importlib.util.spec_from_file_location(f"{package_name}.api.job_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)
|
|
|
|
if created_aiohttp_stub:
|
|
sys.modules.pop("aiohttp", None)
|
|
if created_torch_stub:
|
|
sys.modules.pop("torch", None)
|
|
if created_pil_stub:
|
|
sys.modules.pop("PIL.Image", None)
|
|
sys.modules.pop("PIL", None)
|
|
|
|
return module
|
|
|
|
|
|
job_routes = _load_job_routes_module()
|
|
|
|
|
|
class DistributedQueueEndpointTests(unittest.IsolatedAsyncioTestCase):
|
|
async def test_distributed_queue_happy_path_returns_prompt_id(self):
|
|
request = _FakeRequest(
|
|
{
|
|
"prompt": {"1": {"class_type": "Node"}},
|
|
"enabled_worker_ids": ["w1"],
|
|
"client_id": "client-1",
|
|
"auto_prepare": True,
|
|
}
|
|
)
|
|
with patch.object(
|
|
job_routes,
|
|
"orchestrate_distributed_execution",
|
|
new=AsyncMock(return_value=("prompt_123", 2)),
|
|
):
|
|
response = await job_routes.distributed_queue_endpoint(request)
|
|
|
|
self.assertEqual(response.status, 200)
|
|
self.assertEqual(response.payload.get("prompt_id"), "prompt_123")
|
|
self.assertTrue(response.payload.get("auto_prepare_supported"))
|
|
|
|
async def test_distributed_queue_missing_prompt_returns_400(self):
|
|
request = _FakeRequest(
|
|
{
|
|
"enabled_worker_ids": ["w1"],
|
|
"client_id": "client-1",
|
|
}
|
|
)
|
|
response = await job_routes.distributed_queue_endpoint(request)
|
|
self.assertEqual(response.status, 400)
|
|
self.assertIn("prompt", response.payload.get("message", "").lower())
|
|
|
|
async def test_distributed_queue_missing_enabled_worker_ids_returns_400(self):
|
|
request = _FakeRequest(
|
|
{
|
|
"prompt": {"1": {"class_type": "Node"}},
|
|
"client_id": "client-1",
|
|
}
|
|
)
|
|
response = await job_routes.distributed_queue_endpoint(request)
|
|
self.assertEqual(response.status, 400)
|
|
self.assertIn("enabled_worker_ids", response.payload.get("message", "").lower())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|