import asyncio import importlib.util import io import sys import types import unittest from dataclasses import dataclass, field from pathlib import Path from PIL import Image from tests.api.harness import ( bootstrap_test_package, cleanup_optional_module, install_aiohttp_stub, install_request_guards_stub, install_server_stub, ) class _FakeResponse: def __init__(self, payload, status=200): self.payload = payload self.status = status class _FakeRequest: def __init__(self, json_payload=None, post_payload=None, headers=None, query=None): self._json_payload = json_payload self._post_payload = post_payload or {} self.headers = headers or {} self.query = query or {} async def json(self): return self._json_payload async def post(self): return self._post_payload def _load_usdu_routes_module(): module_path = Path(__file__).resolve().parents[2] / "api" / "usdu_routes.py" package_name = "dist_api_usdu_testpkg" bootstrap_test_package(package_name, with_api=True, with_utils=True, with_upscale=True) schemas_module = types.ModuleType(f"{package_name}.api.schemas") def _require_bool_literal(value, field_name: str = "value"): if isinstance(value, bool): return value if isinstance(value, str): normalized = value.strip().lower() if normalized == "true": return True if normalized == "false": return False raise ValueError(f"Field '{field_name}' must be a boolean literal.") schemas_module.require_bool_literal = _require_bool_literal schemas_module.is_authorized_request = lambda _request, _config: True sys.modules[f"{package_name}.api.schemas"] = schemas_module install_request_guards_stub(package_name) prompt_server_holder = { "value": types.SimpleNamespace( distributed_tile_jobs_lock=asyncio.Lock(), distributed_pending_tile_jobs={}, ) } created_aiohttp_stub = install_aiohttp_stub( lambda payload, status=200: _FakeResponse(payload, status=status) ) install_server_stub() logging_module = types.ModuleType(f"{package_name}.utils.logging") logging_module.debug_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: {"settings": {}} sys.modules[f"{package_name}.utils.config"] = config_module usdu_management_module = types.ModuleType(f"{package_name}.utils.usdu_management") usdu_management_module.MAX_PAYLOAD_SIZE = 1024 sys.modules[f"{package_name}.utils.usdu_management"] = usdu_management_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 job_store_module = types.ModuleType(f"{package_name}.upscale.job_store") job_store_module.MAX_PAYLOAD_SIZE = 1024 job_store_module.ensure_tile_jobs_initialized = lambda: prompt_server_holder["value"] async def _init_dynamic_job(multi_job_id, batch_size, enabled_workers, all_indices=None): prompt_server = prompt_server_holder["value"] async with prompt_server.distributed_tile_jobs_lock: if multi_job_id in prompt_server.distributed_pending_tile_jobs: return job_data = ImageJobState(multi_job_id=multi_job_id) job_data.worker_status = {str(worker_id): 0 for worker_id in (enabled_workers or [])} job_data.assigned_to_workers = {str(worker_id): [] for worker_id in (enabled_workers or [])} indices = all_indices if all_indices is not None else list(range(int(batch_size or 0))) for idx in indices: await job_data.pending_images.put(int(idx)) prompt_server.distributed_pending_tile_jobs[multi_job_id] = job_data job_store_module.init_dynamic_job = _init_dynamic_job sys.modules[f"{package_name}.upscale.job_store"] = job_store_module job_models_module = types.ModuleType(f"{package_name}.upscale.job_models") class BaseJobState: pass @dataclass class TileJobState(BaseJobState): multi_job_id: str mode: str = field(default="static", init=False) queue: asyncio.Queue = field(default_factory=asyncio.Queue) pending_tasks: asyncio.Queue = field(default_factory=asyncio.Queue) completed_tasks: dict = field(default_factory=dict) worker_status: dict = field(default_factory=dict) assigned_to_workers: dict = field(default_factory=dict) batch_size: int = 0 num_tiles_per_image: int = 0 batched_static: bool = False @dataclass class ImageJobState(BaseJobState): multi_job_id: str mode: str = field(default="dynamic", init=False) queue: asyncio.Queue = field(default_factory=asyncio.Queue) pending_images: asyncio.Queue = field(default_factory=asyncio.Queue) completed_images: dict = field(default_factory=dict) worker_status: dict = field(default_factory=dict) assigned_to_workers: dict = field(default_factory=dict) batch_size: int = 0 num_tiles_per_image: int = 0 batched_static: bool = False @property def pending_tasks(self): return self.pending_images @property def completed_tasks(self): return self.completed_images job_models_module.BaseJobState = BaseJobState job_models_module.TileJobState = TileJobState job_models_module.ImageJobState = ImageJobState sys.modules[f"{package_name}.upscale.job_models"] = job_models_module parsers_module = types.ModuleType(f"{package_name}.upscale.payload_parsers") parsers_module.parse_tiles_from_form = lambda _data: [] parsers_module._parse_tiles_from_form = lambda _data: [] sys.modules[f"{package_name}.upscale.payload_parsers"] = parsers_module spec = importlib.util.spec_from_file_location(f"{package_name}.api.usdu_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) cleanup_optional_module("aiohttp", created_aiohttp_stub) module.web = types.SimpleNamespace( json_response=lambda payload, status=200: _FakeResponse(payload, status=status) ) module._prompt_server_holder = prompt_server_holder module._TileJobState = TileJobState module._ImageJobState = ImageJobState return module usdu_routes = _load_usdu_routes_module() class _UploadField: def __init__(self, data): self.file = io.BytesIO(data) def _tiny_png_bytes(): image = Image.new("RGB", (1, 1), (255, 0, 0)) buf = io.BytesIO() image.save(buf, format="PNG") return buf.getvalue() class USDURoutesTests(unittest.IsolatedAsyncioTestCase): async def asyncSetUp(self): usdu_routes._prompt_server_holder["value"] = types.SimpleNamespace( distributed_tile_jobs_lock=asyncio.Lock(), distributed_pending_tile_jobs={}, ) async def test_heartbeat_updates_worker_status(self): prompt_server = usdu_routes._prompt_server_holder["value"] job_data = usdu_routes._TileJobState("job-1") prompt_server.distributed_pending_tile_jobs["job-1"] = job_data request = _FakeRequest(json_payload={"worker_id": "worker-a", "multi_job_id": "job-1"}) response = await usdu_routes.heartbeat_endpoint(request) self.assertEqual(response.status, 200) self.assertEqual(response.payload.get("status"), "success") self.assertIn("worker-a", job_data.worker_status) async def test_heartbeat_missing_fields_returns_400(self): request = _FakeRequest(json_payload={"worker_id": "worker-a"}) response = await usdu_routes.heartbeat_endpoint(request) self.assertEqual(response.status, 400) self.assertIn("missing", response.payload.get("message", "").lower()) async def test_request_image_dynamic_assigns_next_image(self): prompt_server = usdu_routes._prompt_server_holder["value"] job_data = usdu_routes._ImageJobState("job-2") await job_data.pending_images.put(7) prompt_server.distributed_pending_tile_jobs["job-2"] = job_data request = _FakeRequest(json_payload={"worker_id": "worker-a", "multi_job_id": "job-2"}) response = await usdu_routes.request_image_endpoint(request) self.assertEqual(response.status, 200) self.assertEqual(response.payload.get("kind"), "image") self.assertEqual(response.payload.get("task_idx"), 7) self.assertEqual(response.payload.get("estimated_remaining"), 0) self.assertEqual(job_data.assigned_to_workers["worker-a"], [7]) self.assertIn("worker-a", job_data.worker_status) async def test_request_image_static_assigns_tile_and_batched_flag(self): prompt_server = usdu_routes._prompt_server_holder["value"] job_data = usdu_routes._TileJobState("job-3") job_data.batched_static = True await job_data.pending_tasks.put(4) prompt_server.distributed_pending_tile_jobs["job-3"] = job_data request = _FakeRequest(json_payload={"worker_id": "worker-a", "multi_job_id": "job-3"}) response = await usdu_routes.request_image_endpoint(request) self.assertEqual(response.status, 200) self.assertEqual(response.payload.get("kind"), "tile") self.assertEqual(response.payload.get("task_idx"), 4) self.assertTrue(response.payload.get("batched_static")) self.assertEqual(job_data.assigned_to_workers["worker-a"], [4]) async def test_submit_tiles_completion_signal_enqueues_last_marker(self): prompt_server = usdu_routes._prompt_server_holder["value"] job_data = usdu_routes._TileJobState("job-4") prompt_server.distributed_pending_tile_jobs["job-4"] = job_data request = _FakeRequest( post_payload={ "multi_job_id": "job-4", "worker_id": "worker-a", "batch_size": "0", "is_last": "true", }, headers={"content-length": "128"}, ) response = await usdu_routes.submit_tiles_endpoint(request) self.assertEqual(response.status, 200) queued = await job_data.queue.get() self.assertEqual(queued["worker_id"], "worker-a") self.assertTrue(queued["is_last"]) self.assertEqual(queued["tiles"], []) async def test_submit_image_enqueues_processed_image_payload(self): prompt_server = usdu_routes._prompt_server_holder["value"] job_data = usdu_routes._ImageJobState("job-5") prompt_server.distributed_pending_tile_jobs["job-5"] = job_data request = _FakeRequest( post_payload={ "multi_job_id": "job-5", "worker_id": "worker-a", "image_idx": "2", "full_image": _UploadField(_tiny_png_bytes()), "is_last": "false", }, headers={"content-length": "256"}, ) response = await usdu_routes.submit_image_endpoint(request) self.assertEqual(response.status, 200) queued = await job_data.queue.get() self.assertEqual(queued["worker_id"], "worker-a") self.assertEqual(queued["image_idx"], 2) self.assertIn("image", queued) async def test_job_status_endpoint_reports_ready(self): prompt_server = usdu_routes._prompt_server_holder["value"] prompt_server.distributed_pending_tile_jobs["job-6"] = usdu_routes._TileJobState("job-6") request = _FakeRequest(query={"multi_job_id": "job-6"}) response = await usdu_routes.job_status_endpoint(request) self.assertEqual(response.status, 200) self.assertTrue(response.payload.get("ready")) if __name__ == "__main__": unittest.main()