chore: commit all pending distributed changes
This commit is contained in:
Robert Wojciechowski
@@ -34,6 +34,10 @@ def convert_paths_for_platform(obj, target_separator):
|
||||
is_absolute = trimmed.startswith("/") or trimmed.startswith("\\\\")
|
||||
has_protocol = bool(re.match(r"^\w+://", trimmed))
|
||||
|
||||
# URLs are not local paths and should never be separator-normalized.
|
||||
if has_protocol:
|
||||
return trimmed
|
||||
|
||||
# Keep relative media-style paths in forward-slash form (Comfy-style annotated paths).
|
||||
if not has_drive and not is_absolute and not has_protocol and IMAGE_OR_VIDEO_RE.search(trimmed):
|
||||
return re.sub(r"[\\]+", "/", trimmed)
|
||||
|
||||
@@ -140,7 +140,8 @@ def prune_prompt_for_worker(prompt_obj):
|
||||
if node is not None:
|
||||
pruned_prompt[node_id] = json.loads(json.dumps(node))
|
||||
|
||||
next_id = _create_numeric_id_generator(pruned_prompt)
|
||||
# Generate IDs from the original prompt so we never reuse IDs from pruned downstream nodes.
|
||||
next_id = _create_numeric_id_generator(prompt_obj)
|
||||
for dist_id in distributed_ids:
|
||||
if dist_id not in pruned_prompt:
|
||||
continue
|
||||
@@ -187,7 +188,8 @@ def prepare_delegate_master_prompt(prompt_obj, collector_ids):
|
||||
f"Removed upstream reference '{input_name}' from node {node_id} for delegate-only master prompt."
|
||||
)
|
||||
|
||||
next_id = _create_numeric_id_generator(pruned_prompt)
|
||||
# Generate IDs from the original prompt to avoid ID collisions with pruned nodes.
|
||||
next_id = _create_numeric_id_generator(prompt_obj)
|
||||
for collector_id in collector_ids:
|
||||
collector_entry = pruned_prompt.get(collector_id)
|
||||
if not collector_entry:
|
||||
|
||||
+31
@@ -0,0 +1,31 @@
|
||||
# conftest.py — project-level pytest configuration.
|
||||
#
|
||||
# Problem: custom_nodes/ComfyUI-Distributed/__init__.py uses relative imports
|
||||
# (from .distributed import ...) that fail when pytest tries to import it as a
|
||||
# standalone module during Package.setup() for the root package node.
|
||||
#
|
||||
# Fix: patch Package.setup() to skip the root-package's __init__.py import.
|
||||
# All actual package context is provided by each test module via
|
||||
# importlib.util.spec_from_file_location with synthetic stub packages.
|
||||
|
||||
from _pytest.python import Package
|
||||
|
||||
_orig_pkg_setup = Package.setup
|
||||
|
||||
|
||||
def _patched_pkg_setup(self) -> None:
|
||||
# Skip the root package setup — its __init__.py uses relative imports
|
||||
# that require a parent package (ComfyUI's plugin loader) which is not
|
||||
# available in the test environment.
|
||||
if self.path == self.config.rootpath:
|
||||
return
|
||||
_orig_pkg_setup(self)
|
||||
|
||||
|
||||
Package.setup = _patched_pkg_setup
|
||||
|
||||
collect_ignore = [
|
||||
"__init__.py",
|
||||
"distributed.py",
|
||||
"distributed_upscale.py",
|
||||
]
|
||||
File diff suppressed because it is too large
Load Diff
Generated
+1419
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,12 @@
|
||||
{
|
||||
"name": "comfyui-distributed-web-tests",
|
||||
"private": true,
|
||||
"type": "module",
|
||||
"scripts": {
|
||||
"test:web": "bash ./scripts/test-web.sh",
|
||||
"test:web:watch": "bash ./scripts/test-web.sh --watch"
|
||||
},
|
||||
"devDependencies": {
|
||||
"vitest": "^2.1.9"
|
||||
}
|
||||
}
|
||||
@@ -13,3 +13,7 @@ Repository = "https://github.com/robertvoy/ComfyUI-Distributed"
|
||||
PublisherId = "robertvoy"
|
||||
DisplayName = "ComfyUI-Distributed"
|
||||
Icon = "https://raw.githubusercontent.com/robertvoy/ComfyUI-Distributed/refs/heads/main/web/distributed-logo-icon.png"
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
testpaths = ["tests"]
|
||||
pythonpath = ["."]
|
||||
|
||||
Executable
+29
@@ -0,0 +1,29 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
|
||||
REPO_ROOT="$(cd -- "${SCRIPT_DIR}/.." && pwd)"
|
||||
|
||||
export NVM_DIR="${NVM_DIR:-$HOME/.nvm}"
|
||||
if [[ -s "${NVM_DIR}/nvm.sh" ]]; then
|
||||
# shellcheck source=/dev/null
|
||||
. "${NVM_DIR}/nvm.sh"
|
||||
fi
|
||||
|
||||
if ! command -v node >/dev/null 2>&1 || ! command -v npm >/dev/null 2>&1; then
|
||||
echo "[test-web] node/npm are not available." >&2
|
||||
echo "[test-web] Install nvm and Node, or ensure node/npm are on PATH." >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if [[ -f "${REPO_ROOT}/.nvmrc" ]] && command -v nvm >/dev/null 2>&1; then
|
||||
nvm use >/dev/null
|
||||
fi
|
||||
|
||||
cd "${REPO_ROOT}"
|
||||
|
||||
if [[ "${1:-}" == "--watch" ]]; then
|
||||
exec npx vitest web/tests
|
||||
fi
|
||||
|
||||
exec npx vitest run web/tests
|
||||
@@ -1 +0,0 @@
|
||||
|
||||
@@ -74,7 +74,9 @@ def _load_job_routes_module():
|
||||
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,
|
||||
@@ -84,7 +86,9 @@ def _load_job_routes_module():
|
||||
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
|
||||
@@ -165,6 +169,11 @@ def _load_job_routes_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
|
||||
|
||||
|
||||
@@ -0,0 +1,229 @@
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def _load_media_sync_module():
|
||||
module_path = Path(__file__).resolve().parents[2] / "api" / "orchestration" / "media_sync.py"
|
||||
package_name = "dist_ms_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
|
||||
|
||||
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 worker, endpoint="": f"http://localhost{endpoint}"
|
||||
|
||||
async def _fake_session():
|
||||
raise RuntimeError("network calls not used in pure-function tests")
|
||||
|
||||
network_module.get_client_session = _fake_session
|
||||
sys.modules[f"{package_name}.utils.network"] = network_module
|
||||
|
||||
trace_module = types.ModuleType(f"{package_name}.utils.trace_logger")
|
||||
trace_module.trace_debug = lambda *_args, **_kwargs: None
|
||||
trace_module.trace_info = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.trace_logger"] = trace_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):
|
||||
pass
|
||||
|
||||
class _FormData:
|
||||
def add_field(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
aiohttp_module.ClientTimeout = _ClientTimeout
|
||||
aiohttp_module.FormData = _FormData
|
||||
sys.modules["aiohttp"] = aiohttp_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
f"{package_name}.api.orchestration.media_sync",
|
||||
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
|
||||
|
||||
|
||||
ms = _load_media_sync_module()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# convert_paths_for_platform
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class ConvertPathsForPlatformTests(unittest.TestCase):
|
||||
def test_forward_slash_target_normalises_backslashes(self):
|
||||
obj = {"ckpt_name": "C:\\Models\\model.safetensors"}
|
||||
result = ms.convert_paths_for_platform(obj, "/")
|
||||
self.assertEqual(result["ckpt_name"], "C:/Models/model.safetensors")
|
||||
|
||||
def test_backslash_target_normalises_forward_slashes(self):
|
||||
obj = {"ckpt_name": "/models/checkpoints/model.safetensors"}
|
||||
result = ms.convert_paths_for_platform(obj, "\\")
|
||||
self.assertIn("\\", result["ckpt_name"])
|
||||
self.assertNotIn("/", result["ckpt_name"])
|
||||
|
||||
def test_relative_media_paths_always_stay_forward_slash(self):
|
||||
"""Relative image/video paths (Comfy annotated style) must not be backslash-ified."""
|
||||
obj = {"image": "subfolder/my_photo.png"}
|
||||
result = ms.convert_paths_for_platform(obj, "\\")
|
||||
self.assertEqual(result["image"], "subfolder/my_photo.png")
|
||||
|
||||
def test_annotated_relative_media_path_stays_forward_slash(self):
|
||||
obj = {"image": "input/frame.jpg [abc123]"}
|
||||
result = ms.convert_paths_for_platform(obj, "\\")
|
||||
self.assertIn("/", result["image"])
|
||||
self.assertNotIn("\\", result["image"].split("[")[0])
|
||||
|
||||
def test_non_filename_strings_are_untouched(self):
|
||||
obj = {"prompt": "a beautiful cat", "count": 5}
|
||||
result = ms.convert_paths_for_platform(obj, "\\")
|
||||
self.assertEqual(result["prompt"], "a beautiful cat")
|
||||
self.assertEqual(result["count"], 5)
|
||||
|
||||
def test_url_strings_are_untouched(self):
|
||||
obj = {"url": "https://example.com/model.safetensors"}
|
||||
result = ms.convert_paths_for_platform(obj, "\\")
|
||||
self.assertEqual(result["url"], "https://example.com/model.safetensors")
|
||||
|
||||
def test_invalid_separator_returns_obj_unchanged(self):
|
||||
obj = {"ckpt_name": "/models/model.safetensors"}
|
||||
result = ms.convert_paths_for_platform(obj, "|")
|
||||
self.assertEqual(result, obj)
|
||||
|
||||
def test_nested_dict_is_processed_recursively(self):
|
||||
obj = {"node": {"ckpt_name": "C:\\Models\\model.safetensors"}}
|
||||
result = ms.convert_paths_for_platform(obj, "/")
|
||||
self.assertEqual(result["node"]["ckpt_name"], "C:/Models/model.safetensors")
|
||||
|
||||
def test_list_items_are_processed_recursively(self):
|
||||
obj = [{"ckpt_name": "C:\\Models\\model.safetensors"}, "plain string"]
|
||||
result = ms.convert_paths_for_platform(obj, "/")
|
||||
self.assertEqual(result[0]["ckpt_name"], "C:/Models/model.safetensors")
|
||||
self.assertEqual(result[1], "plain string")
|
||||
|
||||
def test_non_string_scalar_values_are_untouched(self):
|
||||
obj = {"seed": 42, "enabled": True, "ratio": 1.5}
|
||||
result = ms.convert_paths_for_platform(obj, "/")
|
||||
self.assertEqual(result["seed"], 42)
|
||||
self.assertTrue(result["enabled"])
|
||||
|
||||
def test_absolute_unix_path_to_windows(self):
|
||||
obj = {"lora": "/home/user/loras/my_lora.safetensors"}
|
||||
result = ms.convert_paths_for_platform(obj, "\\")
|
||||
self.assertNotIn("/", result["lora"])
|
||||
|
||||
def test_already_normalised_path_is_idempotent(self):
|
||||
obj = {"ckpt": "C:/Models/model.safetensors"}
|
||||
result = ms.convert_paths_for_platform(obj, "/")
|
||||
self.assertEqual(result["ckpt"], "C:/Models/model.safetensors")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _find_media_references
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class FindMediaReferencesTests(unittest.TestCase):
|
||||
def test_finds_image_input(self):
|
||||
prompt = {"1": {"class_type": "LoadImage", "inputs": {"image": "photo.png"}}}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertIn("photo.png", refs)
|
||||
|
||||
def test_finds_video_input(self):
|
||||
prompt = {"1": {"class_type": "LoadVideo", "inputs": {"video": "clip.mp4"}}}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertIn("clip.mp4", refs)
|
||||
|
||||
def test_strips_annotation_suffix(self):
|
||||
prompt = {"1": {"class_type": "LoadImage", "inputs": {"image": "photo.jpg [abc123]"}}}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertIn("photo.jpg", refs)
|
||||
self.assertFalse(any("[" in r for r in refs))
|
||||
|
||||
def test_normalises_backslashes_in_path(self):
|
||||
prompt = {"1": {"class_type": "LoadImage", "inputs": {"image": "sub\\img.png"}}}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertIn("sub/img.png", refs)
|
||||
|
||||
def test_ignores_non_media_text_inputs(self):
|
||||
prompt = {"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "a cat"}}}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertEqual(refs, [])
|
||||
|
||||
def test_ignores_node_link_values(self):
|
||||
"""Inputs that are [node_id, slot] lists should be ignored."""
|
||||
prompt = {"1": {"class_type": "Anything", "inputs": {"image": ["2", 0]}}}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertEqual(refs, [])
|
||||
|
||||
def test_deduplicates_same_file_across_nodes(self):
|
||||
prompt = {
|
||||
"1": {"class_type": "LoadImage", "inputs": {"image": "cat.png"}},
|
||||
"2": {"class_type": "LoadImage", "inputs": {"image": "cat.png"}},
|
||||
}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertEqual(len(refs), 1)
|
||||
|
||||
def test_returns_sorted_list(self):
|
||||
prompt = {
|
||||
"1": {"class_type": "LoadImage", "inputs": {"image": "z_image.png"}},
|
||||
"2": {"class_type": "LoadImage", "inputs": {"image": "a_image.jpg"}},
|
||||
}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertEqual(refs, sorted(refs))
|
||||
|
||||
def test_ignores_non_dict_nodes(self):
|
||||
prompt = {"1": "not a node dict", "2": {"class_type": "LoadImage", "inputs": {"image": "img.png"}}}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertIn("img.png", refs)
|
||||
|
||||
def test_empty_prompt_returns_empty_list(self):
|
||||
self.assertEqual(ms._find_media_references({}), [])
|
||||
|
||||
def test_multiple_media_types_all_found(self):
|
||||
prompt = {
|
||||
"1": {"class_type": "LoadImage", "inputs": {"image": "frame.png"}},
|
||||
"2": {"class_type": "LoadVideo", "inputs": {"video": "clip.mp4"}},
|
||||
}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertIn("frame.png", refs)
|
||||
self.assertIn("clip.mp4", refs)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,304 @@
|
||||
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
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
class _Routes:
|
||||
def post(self, _path):
|
||||
def _decorator(fn):
|
||||
return fn
|
||||
|
||||
return _decorator
|
||||
|
||||
def get(self, _path):
|
||||
def _decorator(fn):
|
||||
return fn
|
||||
|
||||
return _decorator
|
||||
|
||||
|
||||
def _load_usdu_routes_module():
|
||||
module_path = Path(__file__).resolve().parents[2] / "api" / "usdu_routes.py"
|
||||
package_name = "dist_api_usdu_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
|
||||
|
||||
api_pkg = types.ModuleType(f"{package_name}.api")
|
||||
api_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.api"] = api_pkg
|
||||
|
||||
upscale_pkg = types.ModuleType(f"{package_name}.upscale")
|
||||
upscale_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.upscale"] = upscale_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
prompt_server_holder = {
|
||||
"value": types.SimpleNamespace(
|
||||
distributed_tile_jobs_lock=asyncio.Lock(),
|
||||
distributed_pending_tile_jobs={},
|
||||
)
|
||||
}
|
||||
|
||||
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 = types.ModuleType("server")
|
||||
server_module.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace(routes=_Routes()))
|
||||
sys.modules["server"] = server_module
|
||||
|
||||
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
|
||||
|
||||
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"]
|
||||
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: []
|
||||
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)
|
||||
|
||||
if created_aiohttp_stub:
|
||||
sys.modules.pop("aiohttp", None)
|
||||
|
||||
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("image_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("tile_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()
|
||||
@@ -143,7 +143,9 @@ def _load_worker_routes_module():
|
||||
server_module.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace(routes=_Routes()))
|
||||
sys.modules["server"] = server_module
|
||||
|
||||
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,
|
||||
@@ -219,6 +221,8 @@ def _load_worker_routes_module():
|
||||
|
||||
if created_aiohttp_stub:
|
||||
sys.modules.pop("aiohttp", None)
|
||||
if created_torch_stub:
|
||||
sys.modules.pop("torch", None)
|
||||
|
||||
return module
|
||||
|
||||
|
||||
@@ -0,0 +1,2 @@
|
||||
# This conftest.py marks the tests/ directory as the pytest collection root,
|
||||
# preventing pytest from traversing into the parent package's __init__.py.
|
||||
@@ -0,0 +1,293 @@
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
|
||||
def _load_config_module():
|
||||
module_path = Path(__file__).resolve().parents[1] / "utils" / "config.py"
|
||||
package_name = "dist_cfg_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
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.logging")
|
||||
logging_module.log = lambda *_args, **_kwargs: None
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.logging"] = logging_module
|
||||
|
||||
constants_module = types.ModuleType(f"{package_name}.constants")
|
||||
constants_module.HEARTBEAT_TIMEOUT = 30
|
||||
sys.modules[f"{package_name}.constants"] = constants_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(f"{package_name}.config", 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
|
||||
|
||||
|
||||
config = _load_config_module()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _merge_with_defaults
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class MergeWithDefaultsTests(unittest.TestCase):
|
||||
def test_non_dict_input_returns_defaults(self):
|
||||
result = config._merge_with_defaults("not a dict", {"key": "default"})
|
||||
self.assertEqual(result, {"key": "default"})
|
||||
|
||||
def test_fills_missing_keys_with_defaults(self):
|
||||
result = config._merge_with_defaults({}, {"a": 1, "b": 2})
|
||||
self.assertEqual(result, {"a": 1, "b": 2})
|
||||
|
||||
def test_loaded_value_overrides_default(self):
|
||||
result = config._merge_with_defaults({"a": 99}, {"a": 1, "b": 2})
|
||||
self.assertEqual(result["a"], 99)
|
||||
self.assertEqual(result["b"], 2)
|
||||
|
||||
def test_nested_dict_merges_recursively(self):
|
||||
defaults = {"settings": {"debug": False, "count": 5}}
|
||||
loaded = {"settings": {"debug": True}}
|
||||
result = config._merge_with_defaults(loaded, defaults)
|
||||
self.assertTrue(result["settings"]["debug"])
|
||||
self.assertEqual(result["settings"]["count"], 5)
|
||||
|
||||
def test_preserves_unknown_keys_for_forward_compatibility(self):
|
||||
result = config._merge_with_defaults({"extra_key": "extra"}, {"a": 1})
|
||||
self.assertEqual(result["extra_key"], "extra")
|
||||
|
||||
def test_none_loaded_value_overrides_default(self):
|
||||
"""Explicitly set None in config should override non-None default."""
|
||||
result = config._merge_with_defaults({"a": None}, {"a": "default"})
|
||||
self.assertIsNone(result["a"])
|
||||
|
||||
def test_non_dict_nested_loaded_value_replaces_dict_default(self):
|
||||
"""If loaded has a scalar where default has a dict, use the scalar."""
|
||||
defaults = {"settings": {"debug": False}}
|
||||
loaded = {"settings": "flat_string"}
|
||||
result = config._merge_with_defaults(loaded, defaults)
|
||||
self.assertEqual(result["settings"], "flat_string")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# load_config
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class LoadConfigTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
config.invalidate_config_cache()
|
||||
|
||||
def tearDown(self):
|
||||
config.invalidate_config_cache()
|
||||
|
||||
def test_returns_defaults_when_file_missing(self):
|
||||
with patch.object(config, "CONFIG_FILE", "/nonexistent/path/config.json"):
|
||||
cfg = config.load_config()
|
||||
defaults = config.get_default_config()
|
||||
self.assertEqual(cfg["settings"]["debug"], defaults["settings"]["debug"])
|
||||
self.assertIn("workers", cfg)
|
||||
|
||||
def test_loads_valid_json_file(self):
|
||||
data = {
|
||||
"workers": [{"id": "w1"}],
|
||||
"master": {"host": "test.host"},
|
||||
"settings": {},
|
||||
"tunnel": {},
|
||||
}
|
||||
with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False, encoding="utf-8") as f:
|
||||
json.dump(data, f)
|
||||
tmp_path = f.name
|
||||
try:
|
||||
with patch.object(config, "CONFIG_FILE", tmp_path):
|
||||
cfg = config.load_config()
|
||||
self.assertEqual(cfg["master"]["host"], "test.host")
|
||||
self.assertEqual(len(cfg["workers"]), 1)
|
||||
finally:
|
||||
os.unlink(tmp_path)
|
||||
|
||||
def test_merges_loaded_file_with_defaults(self):
|
||||
"""Loaded file with partial settings should be filled in from defaults."""
|
||||
data = {"master": {"host": "h"}, "workers": [], "settings": {"debug": True}, "tunnel": {}}
|
||||
with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False, encoding="utf-8") as f:
|
||||
json.dump(data, f)
|
||||
tmp_path = f.name
|
||||
try:
|
||||
with patch.object(config, "CONFIG_FILE", tmp_path):
|
||||
cfg = config.load_config()
|
||||
# debug was set to True
|
||||
self.assertTrue(cfg["settings"]["debug"])
|
||||
# auto_launch_workers is a default key and should be present
|
||||
self.assertIn("auto_launch_workers", cfg["settings"])
|
||||
finally:
|
||||
os.unlink(tmp_path)
|
||||
|
||||
def test_falls_back_to_defaults_on_invalid_json(self):
|
||||
with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False, encoding="utf-8") as f:
|
||||
f.write("{invalid json{{")
|
||||
tmp_path = f.name
|
||||
try:
|
||||
with patch.object(config, "CONFIG_FILE", tmp_path):
|
||||
cfg = config.load_config()
|
||||
self.assertIn("settings", cfg)
|
||||
self.assertIn("workers", cfg)
|
||||
finally:
|
||||
os.unlink(tmp_path)
|
||||
|
||||
def test_second_call_returns_cached_object(self):
|
||||
with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False, encoding="utf-8") as f:
|
||||
json.dump(config.get_default_config(), f)
|
||||
tmp_path = f.name
|
||||
try:
|
||||
with patch.object(config, "CONFIG_FILE", tmp_path):
|
||||
cfg1 = config.load_config()
|
||||
cfg2 = config.load_config()
|
||||
self.assertIs(cfg1, cfg2)
|
||||
finally:
|
||||
os.unlink(tmp_path)
|
||||
|
||||
def test_invalidate_cache_forces_reload(self):
|
||||
data = config.get_default_config()
|
||||
data["master"]["host"] = "first"
|
||||
with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False, encoding="utf-8") as f:
|
||||
json.dump(data, f)
|
||||
tmp_path = f.name
|
||||
try:
|
||||
with patch.object(config, "CONFIG_FILE", tmp_path):
|
||||
cfg1 = config.load_config()
|
||||
config.invalidate_config_cache()
|
||||
data["master"]["host"] = "second"
|
||||
with open(tmp_path, "w", encoding="utf-8") as fh:
|
||||
json.dump(data, fh)
|
||||
cfg2 = config.load_config()
|
||||
self.assertEqual(cfg1["master"]["host"], "first")
|
||||
self.assertEqual(cfg2["master"]["host"], "second")
|
||||
finally:
|
||||
os.unlink(tmp_path)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# save_config
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class SaveConfigTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
config.invalidate_config_cache()
|
||||
|
||||
def tearDown(self):
|
||||
config.invalidate_config_cache()
|
||||
|
||||
def test_saves_and_reloads_correctly(self):
|
||||
data = config.get_default_config()
|
||||
data["master"]["host"] = "saved.host"
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
tmp_path = os.path.join(tmpdir, "config.json")
|
||||
with patch.object(config, "CONFIG_FILE", tmp_path):
|
||||
result = config.save_config(data)
|
||||
self.assertTrue(result)
|
||||
loaded = config.load_config()
|
||||
self.assertEqual(loaded["master"]["host"], "saved.host")
|
||||
|
||||
def test_returns_false_when_path_unwritable(self):
|
||||
with patch.object(config, "CONFIG_FILE", "/nonexistent_dir/config.json"):
|
||||
result = config.save_config({})
|
||||
self.assertFalse(result)
|
||||
|
||||
def test_save_invalidates_cache(self):
|
||||
"""After saving, the cache should be cleared so next load re-reads."""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
tmp_path = os.path.join(tmpdir, "config.json")
|
||||
with patch.object(config, "CONFIG_FILE", tmp_path):
|
||||
data = config.get_default_config()
|
||||
config.save_config(data)
|
||||
# Cache is now None; load_config should re-read
|
||||
self.assertIsNone(config._config_cache)
|
||||
|
||||
def test_written_file_is_valid_json(self):
|
||||
data = config.get_default_config()
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
tmp_path = os.path.join(tmpdir, "config.json")
|
||||
with patch.object(config, "CONFIG_FILE", tmp_path):
|
||||
config.save_config(data)
|
||||
with open(tmp_path, encoding="utf-8") as fh:
|
||||
parsed = json.load(fh)
|
||||
self.assertEqual(parsed["master"], data["master"])
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# get_worker_timeout_seconds
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class GetWorkerTimeoutSecondsTests(unittest.TestCase):
|
||||
def test_returns_configured_value(self):
|
||||
cfg = config.get_default_config()
|
||||
cfg["settings"]["worker_timeout_seconds"] = 120
|
||||
with patch.object(config, "load_config", return_value=cfg):
|
||||
self.assertEqual(config.get_worker_timeout_seconds(), 120)
|
||||
|
||||
def test_clamps_zero_to_one(self):
|
||||
cfg = config.get_default_config()
|
||||
cfg["settings"]["worker_timeout_seconds"] = 0
|
||||
with patch.object(config, "load_config", return_value=cfg):
|
||||
self.assertEqual(config.get_worker_timeout_seconds(), 1)
|
||||
|
||||
def test_clamps_negative_to_one(self):
|
||||
cfg = config.get_default_config()
|
||||
cfg["settings"]["worker_timeout_seconds"] = -10
|
||||
with patch.object(config, "load_config", return_value=cfg):
|
||||
self.assertEqual(config.get_worker_timeout_seconds(), 1)
|
||||
|
||||
def test_falls_back_to_provided_default_when_key_missing(self):
|
||||
cfg = config.get_default_config()
|
||||
# worker_timeout_seconds is not present in default config
|
||||
cfg["settings"].pop("worker_timeout_seconds", None)
|
||||
with patch.object(config, "load_config", return_value=cfg):
|
||||
result = config.get_worker_timeout_seconds(default=45)
|
||||
self.assertEqual(result, 45)
|
||||
|
||||
def test_fallback_also_clamped_to_one(self):
|
||||
cfg = config.get_default_config()
|
||||
cfg["settings"].pop("worker_timeout_seconds", None)
|
||||
with patch.object(config, "load_config", return_value=cfg):
|
||||
result = config.get_worker_timeout_seconds(default=0)
|
||||
self.assertEqual(result, 1)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# is_master_delegate_only
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class IsMasterDelegateOnlyTests(unittest.TestCase):
|
||||
def test_returns_false_by_default(self):
|
||||
cfg = config.get_default_config()
|
||||
with patch.object(config, "load_config", return_value=cfg):
|
||||
self.assertFalse(config.is_master_delegate_only())
|
||||
|
||||
def test_returns_true_when_enabled(self):
|
||||
cfg = config.get_default_config()
|
||||
cfg["settings"]["master_delegate_only"] = True
|
||||
with patch.object(config, "load_config", return_value=cfg):
|
||||
self.assertTrue(config.is_master_delegate_only())
|
||||
|
||||
def test_returns_false_on_exception(self):
|
||||
def _raise():
|
||||
raise RuntimeError("config exploded")
|
||||
|
||||
with patch.object(config, "load_config", side_effect=RuntimeError("boom")):
|
||||
self.assertFalse(config.is_master_delegate_only())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,196 @@
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
|
||||
def _load_detection_module():
|
||||
module_path = Path(__file__).resolve().parents[1] / "workers" / "detection.py"
|
||||
package_name = "dist_det_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
|
||||
|
||||
workers_pkg = types.ModuleType(f"{package_name}.workers")
|
||||
workers_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.workers"] = workers_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
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
network_module.normalize_host = lambda value: value
|
||||
|
||||
async def _fake_session():
|
||||
raise RuntimeError("network calls not used 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):
|
||||
pass
|
||||
|
||||
aiohttp_module.ClientTimeout = _ClientTimeout
|
||||
sys.modules["aiohttp"] = aiohttp_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
f"{package_name}.workers.detection",
|
||||
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
|
||||
|
||||
|
||||
detection = _load_detection_module()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# is_docker_environment
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class IsDockerEnvironmentTests(unittest.TestCase):
|
||||
def test_true_when_dockerenv_file_exists(self):
|
||||
with patch.object(detection.os.path, "exists", return_value=True), \
|
||||
patch.dict(detection.os.environ, {}, clear=True), \
|
||||
patch.object(detection.platform, "node", return_value="my-laptop"):
|
||||
self.assertTrue(detection.is_docker_environment())
|
||||
|
||||
def test_true_when_docker_container_env_var_is_set(self):
|
||||
with patch.object(detection.os.path, "exists", return_value=False), \
|
||||
patch.dict(detection.os.environ, {"DOCKER_CONTAINER": "1"}, clear=True), \
|
||||
patch.object(detection.platform, "node", return_value="my-laptop"):
|
||||
self.assertTrue(detection.is_docker_environment())
|
||||
|
||||
def test_true_when_platform_node_contains_docker(self):
|
||||
with patch.object(detection.os.path, "exists", return_value=False), \
|
||||
patch.dict(detection.os.environ, {}, clear=True), \
|
||||
patch.object(detection.platform, "node", return_value="my-docker-host"):
|
||||
self.assertTrue(detection.is_docker_environment())
|
||||
|
||||
def test_false_when_none_of_the_signals_are_present(self):
|
||||
with patch.object(detection.os.path, "exists", return_value=False), \
|
||||
patch.dict(detection.os.environ, {}, clear=True), \
|
||||
patch.object(detection.platform, "node", return_value="my-laptop"):
|
||||
self.assertFalse(detection.is_docker_environment())
|
||||
|
||||
def test_docker_node_name_is_case_insensitive(self):
|
||||
with patch.object(detection.os.path, "exists", return_value=False), \
|
||||
patch.dict(detection.os.environ, {}, clear=True), \
|
||||
patch.object(detection.platform, "node", return_value="My-Docker-Box"):
|
||||
self.assertTrue(detection.is_docker_environment())
|
||||
|
||||
def test_docker_env_var_empty_string_is_falsy(self):
|
||||
"""An empty DOCKER_CONTAINER env var should NOT trigger docker detection."""
|
||||
with patch.object(detection.os.path, "exists", return_value=False), \
|
||||
patch.dict(detection.os.environ, {"DOCKER_CONTAINER": ""}, clear=True), \
|
||||
patch.object(detection.platform, "node", return_value="my-laptop"):
|
||||
self.assertFalse(detection.is_docker_environment())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# is_runpod_environment
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class IsRunpodEnvironmentTests(unittest.TestCase):
|
||||
def test_true_when_runpod_pod_id_is_set(self):
|
||||
with patch.dict(detection.os.environ, {"RUNPOD_POD_ID": "pod-abc"}, clear=True):
|
||||
self.assertTrue(detection.is_runpod_environment())
|
||||
|
||||
def test_true_when_runpod_api_key_is_set(self):
|
||||
with patch.dict(detection.os.environ, {"RUNPOD_API_KEY": "key-xyz"}, clear=True):
|
||||
self.assertTrue(detection.is_runpod_environment())
|
||||
|
||||
def test_true_when_both_vars_are_set(self):
|
||||
with patch.dict(
|
||||
detection.os.environ,
|
||||
{"RUNPOD_POD_ID": "pod-abc", "RUNPOD_API_KEY": "key-xyz"},
|
||||
clear=True,
|
||||
):
|
||||
self.assertTrue(detection.is_runpod_environment())
|
||||
|
||||
def test_false_when_neither_var_is_set(self):
|
||||
with patch.dict(detection.os.environ, {}, clear=True):
|
||||
self.assertFalse(detection.is_runpod_environment())
|
||||
|
||||
def test_true_when_pod_id_is_empty_string(self):
|
||||
"""is not None check means even empty string counts as detected."""
|
||||
with patch.dict(detection.os.environ, {"RUNPOD_POD_ID": ""}, clear=True):
|
||||
self.assertTrue(detection.is_runpod_environment())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# is_local_worker (synchronous paths only)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class IsLocalWorkerTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_true_for_localhost_host(self):
|
||||
result = await detection.is_local_worker({"host": "localhost", "port": 8188})
|
||||
self.assertTrue(result)
|
||||
|
||||
async def test_true_for_127_0_0_1(self):
|
||||
result = await detection.is_local_worker({"host": "127.0.0.1", "port": 8188})
|
||||
self.assertTrue(result)
|
||||
|
||||
async def test_true_for_0_0_0_0(self):
|
||||
result = await detection.is_local_worker({"host": "0.0.0.0", "port": 8188})
|
||||
self.assertTrue(result)
|
||||
|
||||
async def test_true_when_type_is_local(self):
|
||||
result = await detection.is_local_worker({"type": "local", "host": "remote.example.com"})
|
||||
self.assertTrue(result)
|
||||
|
||||
async def test_false_for_remote_host(self):
|
||||
result = await detection.is_local_worker({"host": "remote.example.com", "port": 8188})
|
||||
self.assertFalse(result)
|
||||
|
||||
async def test_true_when_no_host_key(self):
|
||||
"""Missing host defaults to 'localhost'."""
|
||||
result = await detection.is_local_worker({"port": 8188})
|
||||
self.assertTrue(result)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# get_machine_id
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class GetMachineIdTests(unittest.TestCase):
|
||||
def test_returns_a_string(self):
|
||||
result = detection.get_machine_id()
|
||||
self.assertIsInstance(result, str)
|
||||
|
||||
def test_returns_non_empty_string(self):
|
||||
result = detection.get_machine_id()
|
||||
self.assertTrue(len(result) > 0)
|
||||
|
||||
def test_stable_across_calls(self):
|
||||
r1 = detection.get_machine_id()
|
||||
r2 = detection.get_machine_id()
|
||||
self.assertEqual(r1, r2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,197 @@
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import sys
|
||||
import time
|
||||
import types
|
||||
import unittest
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def _load_job_timeout_module():
|
||||
module_path = Path(__file__).resolve().parents[1] / "upscale" / "job_timeout.py"
|
||||
package_name = "dist_job_timeout_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
|
||||
|
||||
upscale_pkg = types.ModuleType(f"{package_name}.upscale")
|
||||
upscale_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.upscale"] = upscale_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
config_holder = {"value": {"settings": {}, "workers": []}}
|
||||
probe_holder = {"fn": None}
|
||||
prompt_server_holder = {
|
||||
"value": types.SimpleNamespace(
|
||||
distributed_tile_jobs_lock=asyncio.Lock(),
|
||||
distributed_pending_tile_jobs={},
|
||||
)
|
||||
}
|
||||
|
||||
config_module = types.ModuleType(f"{package_name}.utils.config")
|
||||
config_module.load_config = lambda: config_holder["value"]
|
||||
sys.modules[f"{package_name}.utils.config"] = config_module
|
||||
|
||||
constants_module = types.ModuleType(f"{package_name}.utils.constants")
|
||||
constants_module.HEARTBEAT_TIMEOUT = 60
|
||||
sys.modules[f"{package_name}.utils.constants"] = constants_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
|
||||
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
network_module.build_worker_url = lambda worker: f"http://{worker.get('host', '127.0.0.1')}:{worker.get('port', 8188)}"
|
||||
|
||||
async def _probe_worker(url, timeout=2.0):
|
||||
fn = probe_holder["fn"]
|
||||
if fn is None:
|
||||
return None
|
||||
return await fn(url, timeout)
|
||||
|
||||
network_module.probe_worker = _probe_worker
|
||||
sys.modules[f"{package_name}.utils.network"] = network_module
|
||||
|
||||
job_store_module = types.ModuleType(f"{package_name}.upscale.job_store")
|
||||
job_store_module.ensure_tile_jobs_initialized = lambda: prompt_server_holder["value"]
|
||||
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 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
|
||||
|
||||
@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
|
||||
|
||||
job_models_module.BaseJobState = BaseJobState
|
||||
job_models_module.ImageJobState = ImageJobState
|
||||
job_models_module.TileJobState = TileJobState
|
||||
sys.modules[f"{package_name}.upscale.job_models"] = job_models_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(f"{package_name}.upscale.job_timeout", 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)
|
||||
|
||||
module._config_holder = config_holder
|
||||
module._probe_holder = probe_holder
|
||||
module._prompt_server_holder = prompt_server_holder
|
||||
module._ImageJobState = ImageJobState
|
||||
module._TileJobState = TileJobState
|
||||
return module
|
||||
|
||||
|
||||
jt = _load_job_timeout_module()
|
||||
|
||||
|
||||
class JobTimeoutRequeueTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def asyncSetUp(self):
|
||||
jt._prompt_server_holder["value"] = types.SimpleNamespace(
|
||||
distributed_tile_jobs_lock=asyncio.Lock(),
|
||||
distributed_pending_tile_jobs={},
|
||||
)
|
||||
jt._config_holder["value"] = {
|
||||
"settings": {"worker_timeout_seconds": 5},
|
||||
"workers": [{"id": "worker-1", "host": "worker.local", "port": 8188}],
|
||||
}
|
||||
|
||||
async def test_requeues_only_incomplete_dynamic_tasks_for_timed_out_worker(self):
|
||||
async def _offline_probe(_url, _timeout):
|
||||
return None
|
||||
|
||||
jt._probe_holder["fn"] = _offline_probe
|
||||
prompt_server = jt._prompt_server_holder["value"]
|
||||
job_data = jt._ImageJobState("job-1")
|
||||
job_data.worker_status["worker-1"] = time.time() - 60.0
|
||||
job_data.assigned_to_workers["worker-1"] = [0, 1]
|
||||
job_data.completed_images[1] = "done"
|
||||
prompt_server.distributed_pending_tile_jobs["job-1"] = job_data
|
||||
|
||||
requeued = await jt._check_and_requeue_timed_out_workers("job-1", total_tasks=2)
|
||||
|
||||
self.assertEqual(requeued, 1)
|
||||
self.assertEqual(await job_data.pending_images.get(), 0)
|
||||
self.assertNotIn("worker-1", job_data.worker_status)
|
||||
self.assertEqual(job_data.assigned_to_workers["worker-1"], [])
|
||||
|
||||
async def test_busy_probe_graces_worker_and_skips_requeue(self):
|
||||
async def _busy_probe(_url, _timeout):
|
||||
return {"exec_info": {"queue_remaining": 3}}
|
||||
|
||||
jt._probe_holder["fn"] = _busy_probe
|
||||
prompt_server = jt._prompt_server_holder["value"]
|
||||
job_data = jt._ImageJobState("job-2")
|
||||
old_heartbeat = time.time() - 60.0
|
||||
job_data.worker_status["worker-1"] = old_heartbeat
|
||||
job_data.assigned_to_workers["worker-1"] = [0]
|
||||
prompt_server.distributed_pending_tile_jobs["job-2"] = job_data
|
||||
|
||||
requeued = await jt._check_and_requeue_timed_out_workers("job-2", total_tasks=1)
|
||||
|
||||
self.assertEqual(requeued, 0)
|
||||
self.assertIn("worker-1", job_data.worker_status)
|
||||
self.assertGreaterEqual(job_data.worker_status["worker-1"], old_heartbeat)
|
||||
self.assertTrue(job_data.pending_images.empty())
|
||||
|
||||
async def test_completed_dynamic_task_is_not_requeued(self):
|
||||
async def _offline_probe(_url, _timeout):
|
||||
return None
|
||||
|
||||
jt._probe_holder["fn"] = _offline_probe
|
||||
prompt_server = jt._prompt_server_holder["value"]
|
||||
job_data = jt._ImageJobState("job-3")
|
||||
job_data.worker_status["worker-1"] = time.time() - 60.0
|
||||
job_data.assigned_to_workers["worker-1"] = [7]
|
||||
job_data.completed_images[7] = "complete"
|
||||
prompt_server.distributed_pending_tile_jobs["job-3"] = job_data
|
||||
|
||||
requeued = await jt._check_and_requeue_timed_out_workers("job-3", total_tasks=1)
|
||||
|
||||
self.assertEqual(requeued, 0)
|
||||
self.assertTrue(job_data.pending_images.empty())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,208 @@
|
||||
import importlib.util
|
||||
import io
|
||||
import json
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
try:
|
||||
from PIL import Image as PILImage
|
||||
PIL_AVAILABLE = True
|
||||
except ImportError:
|
||||
PIL_AVAILABLE = False
|
||||
|
||||
|
||||
def _load_payload_parsers_module():
|
||||
# payload_parsers.py has no relative imports; only stdlib + PIL
|
||||
module_path = Path(__file__).resolve().parents[1] / "upscale" / "payload_parsers.py"
|
||||
spec = importlib.util.spec_from_file_location("upscale_payload_parsers", 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
|
||||
|
||||
|
||||
if PIL_AVAILABLE:
|
||||
pp = _load_payload_parsers_module()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _make_png_bytes(width=64, height=64, color=(128, 64, 32)):
|
||||
"""Return raw PNG bytes for a solid-colour image."""
|
||||
img = PILImage.new("RGB", (width, height), color=color)
|
||||
buf = io.BytesIO()
|
||||
img.save(buf, format="PNG")
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
class _MockFileField:
|
||||
"""Minimal multipart file-field stub."""
|
||||
|
||||
class _MockFile:
|
||||
def __init__(self, data: bytes):
|
||||
self._buf = io.BytesIO(data)
|
||||
|
||||
def read(self) -> bytes:
|
||||
return self._buf.read()
|
||||
|
||||
def __init__(self, data: bytes):
|
||||
self.file = self._MockFile(data)
|
||||
|
||||
|
||||
def _make_form(n_tiles, *, padding=None, extra_meta=None, image_color=(128, 64, 32)):
|
||||
"""Build a minimal form-data dict with `n_tiles` tile entries."""
|
||||
image_bytes = _make_png_bytes(color=image_color)
|
||||
metadata = []
|
||||
for i in range(n_tiles):
|
||||
entry = {
|
||||
"tile_idx": i,
|
||||
"x": i * 64,
|
||||
"y": 0,
|
||||
"extracted_width": 64,
|
||||
"extracted_height": 64,
|
||||
}
|
||||
if extra_meta and i < len(extra_meta):
|
||||
entry.update(extra_meta[i])
|
||||
metadata.append(entry)
|
||||
|
||||
form = {"tiles_metadata": json.dumps(metadata)}
|
||||
if padding is not None:
|
||||
form["padding"] = str(padding)
|
||||
for i in range(n_tiles):
|
||||
form[f"tile_{i}"] = _MockFileField(image_bytes)
|
||||
return form
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@unittest.skipUnless(PIL_AVAILABLE, "PIL not installed")
|
||||
class ParseTilesFromFormTests(unittest.TestCase):
|
||||
|
||||
# --- happy paths ---
|
||||
|
||||
def test_single_tile_returns_one_entry(self):
|
||||
tiles = pp._parse_tiles_from_form(_make_form(1))
|
||||
self.assertEqual(len(tiles), 1)
|
||||
|
||||
def test_multiple_tiles_all_returned(self):
|
||||
tiles = pp._parse_tiles_from_form(_make_form(3))
|
||||
self.assertEqual(len(tiles), 3)
|
||||
|
||||
def test_tile_image_is_pil_image(self):
|
||||
tiles = pp._parse_tiles_from_form(_make_form(1))
|
||||
self.assertIsInstance(tiles[0]["image"], PILImage.Image)
|
||||
|
||||
def test_tile_metadata_fields_are_parsed(self):
|
||||
tiles = pp._parse_tiles_from_form(_make_form(1))
|
||||
tile = tiles[0]
|
||||
self.assertEqual(tile["tile_idx"], 0)
|
||||
self.assertEqual(tile["x"], 0)
|
||||
self.assertEqual(tile["y"], 0)
|
||||
self.assertEqual(tile["extracted_width"], 64)
|
||||
self.assertEqual(tile["extracted_height"], 64)
|
||||
|
||||
def test_padding_is_parsed_from_form(self):
|
||||
tiles = pp._parse_tiles_from_form(_make_form(1, padding=16))
|
||||
self.assertEqual(tiles[0]["padding"], 16)
|
||||
|
||||
def test_default_padding_is_zero(self):
|
||||
form = _make_form(1)
|
||||
form.pop("padding", None)
|
||||
tiles = pp._parse_tiles_from_form(form)
|
||||
self.assertEqual(tiles[0]["padding"], 0)
|
||||
|
||||
def test_invalid_padding_string_falls_back_to_zero(self):
|
||||
form = _make_form(1)
|
||||
form["padding"] = "not_a_number"
|
||||
tiles = pp._parse_tiles_from_form(form)
|
||||
self.assertEqual(tiles[0]["padding"], 0)
|
||||
|
||||
def test_optional_batch_idx_included_when_present(self):
|
||||
extra = [{"batch_idx": 2}]
|
||||
tiles = pp._parse_tiles_from_form(_make_form(1, extra_meta=extra))
|
||||
self.assertEqual(tiles[0]["batch_idx"], 2)
|
||||
|
||||
def test_optional_global_idx_included_when_present(self):
|
||||
extra = [{"global_idx": 5}]
|
||||
tiles = pp._parse_tiles_from_form(_make_form(1, extra_meta=extra))
|
||||
self.assertEqual(tiles[0]["global_idx"], 5)
|
||||
|
||||
def test_batch_idx_and_global_idx_absent_when_not_in_metadata(self):
|
||||
tiles = pp._parse_tiles_from_form(_make_form(1))
|
||||
self.assertNotIn("batch_idx", tiles[0])
|
||||
self.assertNotIn("global_idx", tiles[0])
|
||||
|
||||
def test_tile_indices_match_metadata_order(self):
|
||||
tiles = pp._parse_tiles_from_form(_make_form(3))
|
||||
for i, tile in enumerate(tiles):
|
||||
self.assertEqual(tile["tile_idx"], i)
|
||||
|
||||
def test_x_coordinates_reflect_metadata(self):
|
||||
tiles = pp._parse_tiles_from_form(_make_form(3))
|
||||
self.assertEqual(tiles[1]["x"], 64)
|
||||
self.assertEqual(tiles[2]["x"], 128)
|
||||
|
||||
# --- error cases ---
|
||||
|
||||
def test_missing_tiles_metadata_raises_value_error(self):
|
||||
with self.assertRaises(ValueError, msg="Missing tiles_metadata"):
|
||||
pp._parse_tiles_from_form({})
|
||||
|
||||
def test_invalid_json_metadata_raises_value_error(self):
|
||||
form = {"tiles_metadata": "{not valid json}"}
|
||||
with self.assertRaises(ValueError):
|
||||
pp._parse_tiles_from_form(form)
|
||||
|
||||
def test_non_list_metadata_raises_value_error(self):
|
||||
form = {"tiles_metadata": json.dumps({"not": "a list"})}
|
||||
with self.assertRaises(ValueError):
|
||||
pp._parse_tiles_from_form(form)
|
||||
|
||||
def test_missing_tile_file_field_raises_value_error(self):
|
||||
form = {
|
||||
"tiles_metadata": json.dumps([{"tile_idx": 0, "x": 0, "y": 0}]),
|
||||
# tile_0 intentionally omitted
|
||||
}
|
||||
with self.assertRaises(ValueError):
|
||||
pp._parse_tiles_from_form(form)
|
||||
|
||||
def test_tile_field_without_file_attr_raises_value_error(self):
|
||||
form = {
|
||||
"tiles_metadata": json.dumps([{"tile_idx": 0, "x": 0, "y": 0}]),
|
||||
"tile_0": "plain string without .file",
|
||||
}
|
||||
with self.assertRaises(ValueError):
|
||||
pp._parse_tiles_from_form(form)
|
||||
|
||||
def test_non_image_bytes_raises_value_error(self):
|
||||
class _BadFileField:
|
||||
class _BadFile:
|
||||
def read(self):
|
||||
return b"this is definitely not image data"
|
||||
file = _BadFile()
|
||||
|
||||
form = {
|
||||
"tiles_metadata": json.dumps([{"tile_idx": 0, "x": 0, "y": 0}]),
|
||||
"tile_0": _BadFileField(),
|
||||
}
|
||||
with self.assertRaises(ValueError):
|
||||
pp._parse_tiles_from_form(form)
|
||||
|
||||
def test_invalid_metadata_value_type_raises_value_error(self):
|
||||
"""Non-integer metadata fields (x, y, etc.) should raise ValueError."""
|
||||
form = {
|
||||
"tiles_metadata": json.dumps([{"tile_idx": 0, "x": "not_int", "y": 0}]),
|
||||
"tile_0": _MockFileField(_make_png_bytes()),
|
||||
}
|
||||
with self.assertRaises(ValueError):
|
||||
pp._parse_tiles_from_form(form)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,468 @@
|
||||
import importlib.util
|
||||
import json
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def _load_prompt_transform_module():
|
||||
module_path = Path(__file__).resolve().parents[1] / "api" / "orchestration" / "prompt_transform.py"
|
||||
package_name = "dist_pt_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
|
||||
|
||||
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
|
||||
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
f"{package_name}.api.orchestration.prompt_transform",
|
||||
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
|
||||
|
||||
|
||||
pt = _load_prompt_transform_module()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _linear_prompt():
|
||||
"""1 → 2 → 3 → 4(DistributedCollector) → 5(SaveImage)"""
|
||||
return {
|
||||
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {}},
|
||||
"2": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["1", 1]}},
|
||||
"3": {"class_type": "KSampler", "inputs": {"model": ["1", 0], "positive": ["2", 0]}},
|
||||
"4": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}},
|
||||
"5": {"class_type": "SaveImage", "inputs": {"images": ["4", 0]}},
|
||||
}
|
||||
|
||||
|
||||
def _collector_only_prompt():
|
||||
"""1(Checkpoint) → 2(DistributedCollector) [no downstream from 2]"""
|
||||
return {
|
||||
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {}},
|
||||
"2": {"class_type": "DistributedCollector", "inputs": {"images": ["1", 0]}},
|
||||
}
|
||||
|
||||
|
||||
def _delegate_prompt():
|
||||
"""1 → 2 → 3(DistributedCollector) → 4(SaveImage)"""
|
||||
return {
|
||||
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {}},
|
||||
"2": {"class_type": "KSampler", "inputs": {"model": ["1", 0]}},
|
||||
"3": {"class_type": "DistributedCollector", "inputs": {"images": ["2", 0]}},
|
||||
"4": {"class_type": "SaveImage", "inputs": {"images": ["3", 0]}},
|
||||
}
|
||||
|
||||
|
||||
def _apply(prompt, participant_id, enabled_worker_ids=None, delegate_master=False):
|
||||
if enabled_worker_ids is None:
|
||||
enabled_worker_ids = ["worker-a", "worker-b"]
|
||||
idx = pt.PromptIndex(prompt)
|
||||
job_id_map = pt.generate_job_id_map(idx, "run")
|
||||
return pt.apply_participant_overrides(
|
||||
prompt,
|
||||
participant_id=participant_id,
|
||||
enabled_worker_ids=enabled_worker_ids,
|
||||
job_id_map=job_id_map,
|
||||
master_url="http://master.example.com",
|
||||
delegate_master=delegate_master,
|
||||
prompt_index=idx,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# PromptIndex
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class PromptIndexTests(unittest.TestCase):
|
||||
def test_nodes_by_class_groups_correctly(self):
|
||||
prompt = {
|
||||
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {}},
|
||||
"2": {"class_type": "DistributedCollector", "inputs": {}},
|
||||
"3": {"class_type": "DistributedCollector", "inputs": {}},
|
||||
}
|
||||
idx = pt.PromptIndex(prompt)
|
||||
self.assertCountEqual(idx.nodes_for_class("DistributedCollector"), ["2", "3"])
|
||||
self.assertEqual(idx.nodes_for_class("CheckpointLoaderSimple"), ["1"])
|
||||
|
||||
def test_nodes_for_class_unknown_returns_empty(self):
|
||||
idx = pt.PromptIndex({"1": {"class_type": "KSampler", "inputs": {}}})
|
||||
self.assertEqual(idx.nodes_for_class("Nonexistent"), [])
|
||||
|
||||
def test_nodes_without_class_type_are_indexed_under_none(self):
|
||||
prompt = {"1": {"inputs": {}}}
|
||||
idx = pt.PromptIndex(prompt)
|
||||
# Should not raise; nodes_for_class with None key or missing class_type
|
||||
self.assertEqual(idx.nodes_for_class("KSampler"), [])
|
||||
|
||||
def test_copy_prompt_is_a_deep_copy(self):
|
||||
prompt = {"1": {"class_type": "KSampler", "inputs": {"seed": 42}}}
|
||||
idx = pt.PromptIndex(prompt)
|
||||
copy = idx.copy_prompt()
|
||||
copy["1"]["inputs"]["seed"] = 999
|
||||
self.assertEqual(prompt["1"]["inputs"]["seed"], 42)
|
||||
|
||||
def test_has_upstream_direct_connection(self):
|
||||
"""Node 4 reads directly from node 3 (KSampler)."""
|
||||
idx = pt.PromptIndex(_linear_prompt())
|
||||
self.assertTrue(idx.has_upstream("4", "KSampler"))
|
||||
|
||||
def test_has_upstream_transitive_connection(self):
|
||||
"""Node 4 → 3 → 2 → 1 (CheckpointLoaderSimple)."""
|
||||
idx = pt.PromptIndex(_linear_prompt())
|
||||
self.assertTrue(idx.has_upstream("4", "CheckpointLoaderSimple"))
|
||||
|
||||
def test_has_upstream_returns_false_when_no_path(self):
|
||||
idx = pt.PromptIndex(_linear_prompt())
|
||||
# CheckpointLoaderSimple has no upstream nodes
|
||||
self.assertFalse(idx.has_upstream("1", "DistributedCollector"))
|
||||
|
||||
def test_has_upstream_result_is_cached(self):
|
||||
idx = pt.PromptIndex(_linear_prompt())
|
||||
r1 = idx.has_upstream("4", "KSampler")
|
||||
r2 = idx.has_upstream("4", "KSampler")
|
||||
self.assertEqual(r1, r2)
|
||||
self.assertIn(("4", "KSampler"), idx._upstream_cache)
|
||||
|
||||
def test_has_upstream_does_not_infinite_loop_on_cycle(self):
|
||||
"""Cyclic references in inputs should not cause infinite recursion."""
|
||||
prompt = {
|
||||
"1": {"class_type": "A", "inputs": {"x": ["2", 0]}},
|
||||
"2": {"class_type": "B", "inputs": {"x": ["1", 0]}},
|
||||
}
|
||||
idx = pt.PromptIndex(prompt)
|
||||
# Should terminate without error
|
||||
result = idx.has_upstream("1", "NonExistent")
|
||||
self.assertFalse(result)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# find_nodes_by_class
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class FindNodesByClassTests(unittest.TestCase):
|
||||
def test_finds_matching_nodes(self):
|
||||
prompt = {
|
||||
"1": {"class_type": "KSampler", "inputs": {}},
|
||||
"2": {"class_type": "DistributedCollector", "inputs": {}},
|
||||
}
|
||||
result = pt.find_nodes_by_class(prompt, "KSampler")
|
||||
self.assertEqual(result, ["1"])
|
||||
|
||||
def test_returns_empty_when_no_match(self):
|
||||
prompt = {"1": {"class_type": "KSampler", "inputs": {}}}
|
||||
self.assertEqual(pt.find_nodes_by_class(prompt, "DistributedCollector"), [])
|
||||
|
||||
def test_skips_non_dict_nodes(self):
|
||||
prompt = {"1": "not a dict", "2": {"class_type": "KSampler", "inputs": {}}}
|
||||
result = pt.find_nodes_by_class(prompt, "KSampler")
|
||||
self.assertEqual(result, ["2"])
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# prune_prompt_for_worker
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class PrunePromptForWorkerTests(unittest.TestCase):
|
||||
def test_no_distributed_nodes_returns_prompt_unchanged(self):
|
||||
prompt = {
|
||||
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {}},
|
||||
"2": {"class_type": "SaveImage", "inputs": {"images": ["1", 0]}},
|
||||
}
|
||||
result = pt.prune_prompt_for_worker(prompt)
|
||||
self.assertCountEqual(result.keys(), ["1", "2"])
|
||||
|
||||
def test_keeps_collector_and_upstream(self):
|
||||
prompt = _linear_prompt()
|
||||
result = pt.prune_prompt_for_worker(prompt)
|
||||
for node_id in ("1", "2", "3", "4"):
|
||||
self.assertIn(node_id, result)
|
||||
|
||||
def test_removes_downstream_of_collector(self):
|
||||
prompt = _linear_prompt()
|
||||
result = pt.prune_prompt_for_worker(prompt)
|
||||
self.assertNotIn("5", result)
|
||||
|
||||
def test_injects_preview_image_when_downstream_exists(self):
|
||||
prompt = _linear_prompt()
|
||||
result = pt.prune_prompt_for_worker(prompt)
|
||||
preview_nodes = [n for n in result.values() if n.get("class_type") == "PreviewImage"]
|
||||
self.assertEqual(len(preview_nodes), 1)
|
||||
self.assertEqual(preview_nodes[0]["inputs"]["images"], ["4", 0])
|
||||
|
||||
def test_no_preview_image_when_no_downstream(self):
|
||||
result = pt.prune_prompt_for_worker(_collector_only_prompt())
|
||||
preview_nodes = [n for n in result.values() if n.get("class_type") == "PreviewImage"]
|
||||
self.assertEqual(len(preview_nodes), 0)
|
||||
|
||||
def test_unrelated_nodes_are_pruned(self):
|
||||
prompt = {
|
||||
"1": {"class_type": "DistributedCollector", "inputs": {}},
|
||||
"2": {"class_type": "UnrelatedNode", "inputs": {}}, # no connection to 1
|
||||
}
|
||||
result = pt.prune_prompt_for_worker(prompt)
|
||||
self.assertIn("1", result)
|
||||
self.assertNotIn("2", result)
|
||||
|
||||
def test_result_is_a_copy_not_same_object(self):
|
||||
prompt = _linear_prompt()
|
||||
result = pt.prune_prompt_for_worker(prompt)
|
||||
# Mutating the result should not affect the original
|
||||
original_keys = set(prompt.keys())
|
||||
result["NEW"] = {"class_type": "Test", "inputs": {}}
|
||||
self.assertEqual(set(prompt.keys()), original_keys)
|
||||
|
||||
def test_upscale_node_is_treated_as_distributed(self):
|
||||
prompt = {
|
||||
"1": {"class_type": "KSampler", "inputs": {}},
|
||||
"2": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"image": ["1", 0]}},
|
||||
"3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0]}},
|
||||
}
|
||||
result = pt.prune_prompt_for_worker(prompt)
|
||||
self.assertIn("1", result)
|
||||
self.assertIn("2", result)
|
||||
self.assertNotIn("3", result)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# prepare_delegate_master_prompt
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class PrepareDelegateMasterPromptTests(unittest.TestCase):
|
||||
def test_keeps_collector_and_downstream(self):
|
||||
prompt = _delegate_prompt()
|
||||
result = pt.prepare_delegate_master_prompt(prompt, ["3"])
|
||||
self.assertIn("3", result)
|
||||
self.assertIn("4", result)
|
||||
self.assertNotIn("1", result)
|
||||
self.assertNotIn("2", result)
|
||||
|
||||
def test_removes_dangling_upstream_refs(self):
|
||||
"""Collector must not retain dangling refs to pruned upstream nodes."""
|
||||
prompt = _delegate_prompt()
|
||||
result = pt.prepare_delegate_master_prompt(prompt, ["3"])
|
||||
collector_inputs = result["3"].get("inputs", {})
|
||||
# Original "images" pointed at node 2, which is pruned.
|
||||
# It should now point at a newly injected placeholder node.
|
||||
self.assertIn("images", collector_inputs)
|
||||
source_id = str(collector_inputs["images"][0])
|
||||
self.assertNotEqual(source_id, "2")
|
||||
self.assertIn(source_id, result)
|
||||
self.assertEqual(result[source_id].get("class_type"), "DistributedEmptyImage")
|
||||
|
||||
def test_injects_empty_image_placeholder(self):
|
||||
prompt = _delegate_prompt()
|
||||
result = pt.prepare_delegate_master_prompt(prompt, ["3"])
|
||||
empty_nodes = [(nid, n) for nid, n in result.items() if n.get("class_type") == "DistributedEmptyImage"]
|
||||
self.assertEqual(len(empty_nodes), 1)
|
||||
placeholder_id = empty_nodes[0][0]
|
||||
self.assertEqual(result["3"]["inputs"]["images"], [placeholder_id, 0])
|
||||
|
||||
def test_one_placeholder_per_collector(self):
|
||||
"""Two collectors → two placeholders."""
|
||||
prompt = {
|
||||
"1": {"class_type": "DistributedCollector", "inputs": {}},
|
||||
"2": {"class_type": "DistributedCollector", "inputs": {}},
|
||||
"3": {"class_type": "SaveImage", "inputs": {"images": ["1", 0]}},
|
||||
}
|
||||
result = pt.prepare_delegate_master_prompt(prompt, ["1", "2"])
|
||||
empty_nodes = [n for n in result.values() if n.get("class_type") == "DistributedEmptyImage"]
|
||||
self.assertEqual(len(empty_nodes), 2)
|
||||
|
||||
def test_result_is_independent_copy(self):
|
||||
prompt = _delegate_prompt()
|
||||
result = pt.prepare_delegate_master_prompt(prompt, ["3"])
|
||||
result["3"]["inputs"]["NEW"] = "injected"
|
||||
# Original should be untouched
|
||||
self.assertNotIn("NEW", prompt["3"].get("inputs", {}))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# generate_job_id_map
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class GenerateJobIdMapTests(unittest.TestCase):
|
||||
def test_maps_collector_nodes(self):
|
||||
prompt = {
|
||||
"1": {"class_type": "DistributedCollector", "inputs": {}},
|
||||
"2": {"class_type": "KSampler", "inputs": {}},
|
||||
}
|
||||
idx = pt.PromptIndex(prompt)
|
||||
job_map = pt.generate_job_id_map(idx, "prefix")
|
||||
self.assertEqual(job_map["1"], "prefix_1")
|
||||
self.assertNotIn("2", job_map)
|
||||
|
||||
def test_maps_upscale_nodes(self):
|
||||
prompt = {
|
||||
"5": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {}},
|
||||
}
|
||||
idx = pt.PromptIndex(prompt)
|
||||
job_map = pt.generate_job_id_map(idx, "run")
|
||||
self.assertEqual(job_map["5"], "run_5")
|
||||
|
||||
def test_empty_prompt_returns_empty_map(self):
|
||||
idx = pt.PromptIndex({})
|
||||
self.assertEqual(pt.generate_job_id_map(idx, "prefix"), {})
|
||||
|
||||
def test_stable_ids_across_calls(self):
|
||||
prompt = {"1": {"class_type": "DistributedCollector", "inputs": {}}}
|
||||
idx = pt.PromptIndex(prompt)
|
||||
m1 = pt.generate_job_id_map(idx, "run")
|
||||
m2 = pt.generate_job_id_map(idx, "run")
|
||||
self.assertEqual(m1, m2)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# apply_participant_overrides – DistributedCollector
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class ApplyOverridesCollectorTests(unittest.TestCase):
|
||||
def _collector_prompt(self):
|
||||
return {"1": {"class_type": "DistributedCollector", "inputs": {}}}
|
||||
|
||||
def test_worker_sets_is_worker_true(self):
|
||||
result = _apply(self._collector_prompt(), "worker-a")
|
||||
self.assertTrue(result["1"]["inputs"]["is_worker"])
|
||||
|
||||
def test_worker_sets_master_url(self):
|
||||
result = _apply(self._collector_prompt(), "worker-a")
|
||||
self.assertEqual(result["1"]["inputs"]["master_url"], "http://master.example.com")
|
||||
|
||||
def test_worker_sets_worker_id(self):
|
||||
result = _apply(self._collector_prompt(), "worker-a")
|
||||
self.assertEqual(result["1"]["inputs"]["worker_id"], "worker-a")
|
||||
|
||||
def test_worker_sets_delegate_only_false(self):
|
||||
result = _apply(self._collector_prompt(), "worker-a")
|
||||
self.assertFalse(result["1"]["inputs"]["delegate_only"])
|
||||
|
||||
def test_master_sets_is_worker_false(self):
|
||||
result = _apply(self._collector_prompt(), "master")
|
||||
self.assertFalse(result["1"]["inputs"]["is_worker"])
|
||||
|
||||
def test_master_clears_stale_master_url(self):
|
||||
prompt = {"1": {"class_type": "DistributedCollector", "inputs": {"master_url": "stale"}}}
|
||||
result = _apply(prompt, "master")
|
||||
self.assertNotIn("master_url", result["1"]["inputs"])
|
||||
|
||||
def test_master_clears_stale_worker_id(self):
|
||||
prompt = {"1": {"class_type": "DistributedCollector", "inputs": {"worker_id": "stale"}}}
|
||||
result = _apply(prompt, "master")
|
||||
self.assertNotIn("worker_id", result["1"]["inputs"])
|
||||
|
||||
def test_master_with_delegate_master_sets_delegate_only_true(self):
|
||||
result = _apply(self._collector_prompt(), "master", delegate_master=True)
|
||||
self.assertTrue(result["1"]["inputs"]["delegate_only"])
|
||||
|
||||
def test_master_without_delegate_master_sets_delegate_only_false(self):
|
||||
result = _apply(self._collector_prompt(), "master", delegate_master=False)
|
||||
self.assertFalse(result["1"]["inputs"]["delegate_only"])
|
||||
|
||||
def test_enabled_worker_ids_serialized_as_json(self):
|
||||
enabled = ["worker-a", "worker-b"]
|
||||
result = _apply(self._collector_prompt(), "master", enabled_worker_ids=enabled)
|
||||
self.assertEqual(result["1"]["inputs"]["enabled_worker_ids"], json.dumps(enabled))
|
||||
|
||||
def test_multi_job_id_is_set_from_job_map(self):
|
||||
prompt = {"1": {"class_type": "DistributedCollector", "inputs": {}}}
|
||||
idx = pt.PromptIndex(prompt)
|
||||
job_id_map = {"1": "run_abc_1"}
|
||||
result = pt.apply_participant_overrides(
|
||||
prompt,
|
||||
participant_id="worker-a",
|
||||
enabled_worker_ids=["worker-a"],
|
||||
job_id_map=job_id_map,
|
||||
master_url="http://master",
|
||||
delegate_master=False,
|
||||
prompt_index=idx,
|
||||
)
|
||||
self.assertEqual(result["1"]["inputs"]["multi_job_id"], "run_abc_1")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# apply_participant_overrides – DistributedSeed
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class ApplyOverridesSeedTests(unittest.TestCase):
|
||||
def _seed_prompt(self):
|
||||
return {"1": {"class_type": "DistributedSeed", "inputs": {}}}
|
||||
|
||||
def test_worker_sets_is_worker_true(self):
|
||||
result = _apply(self._seed_prompt(), "worker-a")
|
||||
self.assertTrue(result["1"]["inputs"]["is_worker"])
|
||||
|
||||
def test_worker_id_reflects_index_in_enabled_list(self):
|
||||
result = _apply(self._seed_prompt(), "worker-b", enabled_worker_ids=["worker-a", "worker-b"])
|
||||
self.assertEqual(result["1"]["inputs"]["worker_id"], "worker_1")
|
||||
|
||||
def test_master_sets_is_worker_false(self):
|
||||
result = _apply(self._seed_prompt(), "master")
|
||||
self.assertFalse(result["1"]["inputs"]["is_worker"])
|
||||
|
||||
def test_master_sets_empty_worker_id(self):
|
||||
result = _apply(self._seed_prompt(), "master")
|
||||
self.assertEqual(result["1"]["inputs"]["worker_id"], "")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# apply_participant_overrides – UltimateSDUpscaleDistributed
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class ApplyOverridesUpscaleTests(unittest.TestCase):
|
||||
def _upscale_prompt(self):
|
||||
return {"1": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {}}}
|
||||
|
||||
def test_worker_sets_is_worker_true(self):
|
||||
result = _apply(self._upscale_prompt(), "worker-a")
|
||||
self.assertTrue(result["1"]["inputs"]["is_worker"])
|
||||
|
||||
def test_worker_sets_master_url_and_worker_id(self):
|
||||
result = _apply(self._upscale_prompt(), "worker-a")
|
||||
self.assertEqual(result["1"]["inputs"]["master_url"], "http://master.example.com")
|
||||
self.assertEqual(result["1"]["inputs"]["worker_id"], "worker-a")
|
||||
|
||||
def test_master_clears_master_url_and_worker_id(self):
|
||||
prompt = {"1": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"master_url": "x", "worker_id": "y"}}}
|
||||
result = _apply(prompt, "master")
|
||||
self.assertNotIn("master_url", result["1"]["inputs"])
|
||||
self.assertNotIn("worker_id", result["1"]["inputs"])
|
||||
|
||||
def test_collector_downstream_of_upscale_gets_pass_through(self):
|
||||
"""A DistributedCollector that is downstream of UltimateSDUpscaleDistributed → pass_through=True."""
|
||||
prompt = {
|
||||
"1": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {}},
|
||||
"2": {"class_type": "DistributedCollector", "inputs": {"images": ["1", 0]}},
|
||||
}
|
||||
result = _apply(prompt, "worker-a", enabled_worker_ids=["worker-a"])
|
||||
self.assertTrue(result["2"]["inputs"].get("pass_through"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,253 @@
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def _load_static_mode_module():
|
||||
module_path = Path(__file__).resolve().parents[1] / "upscale" / "modes" / "static.py"
|
||||
package_name = "dist_static_mode_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
|
||||
|
||||
upscale_pkg = types.ModuleType(f"{package_name}.upscale")
|
||||
upscale_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.upscale"] = upscale_pkg
|
||||
|
||||
modes_pkg = types.ModuleType(f"{package_name}.upscale.modes")
|
||||
modes_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.upscale.modes"] = modes_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
created_comfy_stub = False
|
||||
if "comfy" not in sys.modules:
|
||||
created_comfy_stub = True
|
||||
comfy_module = types.ModuleType("comfy")
|
||||
model_mgmt = types.ModuleType("comfy.model_management")
|
||||
|
||||
class _InterruptProcessingException(Exception):
|
||||
pass
|
||||
|
||||
model_mgmt.processing_interrupted = lambda: False
|
||||
model_mgmt.throw_exception_if_processing_interrupted = lambda: None
|
||||
model_mgmt.InterruptProcessingException = _InterruptProcessingException
|
||||
|
||||
comfy_module.model_management = model_mgmt
|
||||
sys.modules["comfy"] = comfy_module
|
||||
sys.modules["comfy.model_management"] = model_mgmt
|
||||
|
||||
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")
|
||||
from PIL import Image as PILImage
|
||||
import numpy as np
|
||||
|
||||
def _tensor_to_pil(img_tensor, batch_index=0):
|
||||
return PILImage.fromarray((255 * img_tensor[batch_index].cpu().numpy()).astype(np.uint8))
|
||||
|
||||
def _pil_to_tensor(image):
|
||||
arr = np.array(image).astype(np.float32) / 255.0
|
||||
return torch.from_numpy(arr).unsqueeze(0)
|
||||
|
||||
image_module.tensor_to_pil = _tensor_to_pil
|
||||
image_module.pil_to_tensor = _pil_to_tensor
|
||||
sys.modules[f"{package_name}.utils.image"] = image_module
|
||||
|
||||
async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
|
||||
|
||||
def _run_async_in_server_loop(coro, timeout=None):
|
||||
if timeout is not None:
|
||||
return asyncio.run(asyncio.wait_for(coro, timeout=timeout))
|
||||
return asyncio.run(coro)
|
||||
|
||||
async_helpers_module.run_async_in_server_loop = _run_async_in_server_loop
|
||||
sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
|
||||
|
||||
config_module = types.ModuleType(f"{package_name}.utils.config")
|
||||
config_module.get_worker_timeout_seconds = lambda: 60
|
||||
sys.modules[f"{package_name}.utils.config"] = config_module
|
||||
|
||||
constants_module = types.ModuleType(f"{package_name}.utils.constants")
|
||||
constants_module.HEARTBEAT_INTERVAL = 10.0
|
||||
constants_module.JOB_POLL_INTERVAL = 0.0
|
||||
constants_module.JOB_POLL_MAX_ATTEMPTS = 3
|
||||
constants_module.MAX_BATCH = 20
|
||||
constants_module.TILE_SEND_TIMEOUT = 1.0
|
||||
constants_module.TILE_WAIT_TIMEOUT = 1.0
|
||||
sys.modules[f"{package_name}.utils.constants"] = constants_module
|
||||
|
||||
job_store_module = types.ModuleType(f"{package_name}.upscale.job_store")
|
||||
|
||||
async def _noop(*_args, **_kwargs):
|
||||
return None
|
||||
|
||||
job_store_module.ensure_tile_jobs_initialized = lambda: types.SimpleNamespace(
|
||||
distributed_tile_jobs_lock=asyncio.Lock(),
|
||||
distributed_pending_tile_jobs={},
|
||||
)
|
||||
job_store_module.init_static_job_batched = _noop
|
||||
job_store_module._mark_task_completed = _noop
|
||||
job_store_module._cleanup_job = _noop
|
||||
job_store_module._drain_results_queue = _noop
|
||||
job_store_module._get_completed_count = _noop
|
||||
sys.modules[f"{package_name}.upscale.job_store"] = job_store_module
|
||||
|
||||
job_models_module = types.ModuleType(f"{package_name}.upscale.job_models")
|
||||
|
||||
class _TileJobState:
|
||||
pass
|
||||
|
||||
job_models_module.TileJobState = _TileJobState
|
||||
sys.modules[f"{package_name}.upscale.job_models"] = job_models_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(f"{package_name}.upscale.modes.static", 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_comfy_stub:
|
||||
sys.modules.pop("comfy.model_management", None)
|
||||
sys.modules.pop("comfy", None)
|
||||
|
||||
return module
|
||||
|
||||
|
||||
static_mode = _load_static_mode_module()
|
||||
|
||||
|
||||
class _FakeStaticWorker(static_mode.StaticModeMixin):
|
||||
def __init__(self):
|
||||
self.sent_batches = []
|
||||
self.request_calls = 0
|
||||
self.heartbeat_calls = 0
|
||||
self.job_ready = True
|
||||
self.tile_sequence = [(0, 0, True), (None, 0, True)]
|
||||
|
||||
def round_to_multiple(self, value):
|
||||
return value
|
||||
|
||||
def calculate_tiles(self, _width, _height, _tile_width, _tile_height, _force_uniform_tiles):
|
||||
return [(0, 0)]
|
||||
|
||||
def _poll_job_ready(self, *_args, **_kwargs):
|
||||
return self.job_ready
|
||||
|
||||
async def _request_tile_from_master(self, *_args, **_kwargs):
|
||||
self.request_calls += 1
|
||||
return self.tile_sequence.pop(0)
|
||||
|
||||
async def _send_heartbeat_to_master(self, *_args, **_kwargs):
|
||||
self.heartbeat_calls += 1
|
||||
|
||||
async def send_tiles_batch_to_master(
|
||||
self,
|
||||
processed_tiles,
|
||||
_multi_job_id,
|
||||
_master_url,
|
||||
_padding,
|
||||
_worker_id,
|
||||
is_final_flush=False,
|
||||
):
|
||||
self.sent_batches.append(
|
||||
{
|
||||
"tiles": list(processed_tiles),
|
||||
"is_final_flush": bool(is_final_flush),
|
||||
}
|
||||
)
|
||||
|
||||
def _extract_and_process_tile(self, upscaled_image, *_args, **_kwargs):
|
||||
batch_size = upscaled_image.shape[0]
|
||||
processed_batch = torch.zeros((batch_size, 2, 2, 3), dtype=torch.float32)
|
||||
return processed_batch, 0, 0, 2, 2
|
||||
|
||||
|
||||
def _call_worker_static(fake_worker):
|
||||
image = torch.zeros((1, 4, 4, 3), dtype=torch.float32)
|
||||
return fake_worker._process_worker_static_sync(
|
||||
image,
|
||||
model=None,
|
||||
positive=None,
|
||||
negative=None,
|
||||
vae=None,
|
||||
seed=1,
|
||||
steps=1,
|
||||
cfg=1.0,
|
||||
sampler_name="euler",
|
||||
scheduler="normal",
|
||||
denoise=0.5,
|
||||
tile_width=4,
|
||||
tile_height=4,
|
||||
padding=8,
|
||||
mask_blur=4,
|
||||
force_uniform_tiles=True,
|
||||
tiled_decode=False,
|
||||
multi_job_id="job-1",
|
||||
master_url="http://master:8188",
|
||||
worker_id="worker-1",
|
||||
enabled_workers=["worker-1"],
|
||||
)
|
||||
|
||||
|
||||
class StaticModeWorkerFlowTests(unittest.TestCase):
|
||||
def test_worker_static_aborts_when_job_not_ready(self):
|
||||
worker = _FakeStaticWorker()
|
||||
worker.job_ready = False
|
||||
|
||||
result = _call_worker_static(worker)
|
||||
|
||||
self.assertEqual(result[0].shape[0], 1)
|
||||
self.assertEqual(worker.request_calls, 0)
|
||||
self.assertEqual(worker.heartbeat_calls, 0)
|
||||
self.assertEqual(worker.sent_batches, [])
|
||||
|
||||
def test_worker_static_requests_tiles_and_flushes_final_batch(self):
|
||||
worker = _FakeStaticWorker()
|
||||
|
||||
_call_worker_static(worker)
|
||||
|
||||
self.assertEqual(worker.request_calls, 2) # one tile, then sentinel
|
||||
self.assertEqual(worker.heartbeat_calls, 1)
|
||||
self.assertEqual(len(worker.sent_batches), 1)
|
||||
self.assertTrue(worker.sent_batches[0]["is_final_flush"])
|
||||
tiles = worker.sent_batches[0]["tiles"]
|
||||
self.assertEqual(len(tiles), 1)
|
||||
self.assertEqual(tiles[0]["tile_idx"], 0)
|
||||
self.assertEqual(tiles[0]["global_idx"], 0)
|
||||
self.assertEqual(tiles[0]["batch_idx"], 0)
|
||||
|
||||
def test_flush_empty_final_still_sends_completion_signal(self):
|
||||
worker = _FakeStaticWorker()
|
||||
|
||||
returned = worker._flush_tiles_to_master(
|
||||
[],
|
||||
"job-1",
|
||||
"http://master:8188",
|
||||
8,
|
||||
"worker-1",
|
||||
is_final_flush=True,
|
||||
)
|
||||
|
||||
self.assertEqual(returned, [])
|
||||
self.assertEqual(len(worker.sent_batches), 1)
|
||||
self.assertEqual(worker.sent_batches[0]["tiles"], [])
|
||||
self.assertTrue(worker.sent_batches[0]["is_final_flush"])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,8 @@
|
||||
import { defineConfig } from "vitest/config";
|
||||
|
||||
export default defineConfig({
|
||||
test: {
|
||||
include: ["web/tests/**/*.test.js"],
|
||||
environment: "node",
|
||||
},
|
||||
});
|
||||
+1
-1
@@ -192,7 +192,7 @@ class DistributedExtension {
|
||||
if (statusDot) {
|
||||
if (enabled) {
|
||||
// Enabled: Start with checking state and trigger check
|
||||
this.ui.updateStatusDot(workerId, STATUS_COLORS.OFFLINE_RED, "Checking status...", true);
|
||||
this.ui.updateStatusDot(workerId, STATUS_COLORS.OFFLINE_RED, "Checking status...", false);
|
||||
setTimeout(() => checkWorkerStatus(this, worker), TIMEOUTS.STATUS_CHECK_DELAY);
|
||||
} else {
|
||||
// Disabled: Set to gray
|
||||
|
||||
@@ -29,6 +29,16 @@ describe("execution decision helpers", () => {
|
||||
expect(prompt["2"].inputs.skip_dispatch).toBeUndefined();
|
||||
});
|
||||
|
||||
it("markSkipDispatch on empty prompt does not throw", () => {
|
||||
expect(() => markSkipDispatch({})).not.toThrow();
|
||||
});
|
||||
|
||||
it("markSkipDispatch initialises missing inputs object before setting flag", () => {
|
||||
const prompt = { "1": { class_type: NODE_CLASSES.DISTRIBUTED_QUEUE } };
|
||||
markSkipDispatch(prompt);
|
||||
expect(prompt["1"].inputs.skip_dispatch).toBe(true);
|
||||
});
|
||||
|
||||
it("selectLeastBusyWorker round-robins among idle workers", () => {
|
||||
const extension = {};
|
||||
const statuses = [
|
||||
@@ -41,6 +51,27 @@ describe("execution decision helpers", () => {
|
||||
expect(selectLeastBusyWorker(extension, statuses)).toBe(statuses[0]);
|
||||
});
|
||||
|
||||
it("selectLeastBusyWorker with a single idle worker always returns it", () => {
|
||||
const extension = {};
|
||||
const statuses = [{ worker: { id: "w1" }, queueRemaining: 0 }];
|
||||
expect(selectLeastBusyWorker(extension, statuses)).toBe(statuses[0]);
|
||||
expect(selectLeastBusyWorker(extension, statuses)).toBe(statuses[0]);
|
||||
});
|
||||
|
||||
it("selectLeastBusyWorker round-robin wraps around correctly", () => {
|
||||
const extension = {};
|
||||
const statuses = [
|
||||
{ worker: { id: "w1" }, queueRemaining: 0 },
|
||||
{ worker: { id: "w2" }, queueRemaining: 0 },
|
||||
{ worker: { id: "w3" }, queueRemaining: 0 },
|
||||
];
|
||||
const results = Array.from({ length: 4 }, () => selectLeastBusyWorker(extension, statuses));
|
||||
expect(results[0]).toBe(statuses[0]);
|
||||
expect(results[1]).toBe(statuses[1]);
|
||||
expect(results[2]).toBe(statuses[2]);
|
||||
expect(results[3]).toBe(statuses[0]); // wraps around
|
||||
});
|
||||
|
||||
it("selectLeastBusyWorker chooses shortest queue when no idle workers", () => {
|
||||
const extension = {};
|
||||
const statuses = [
|
||||
@@ -52,6 +83,12 @@ describe("execution decision helpers", () => {
|
||||
expect(selectLeastBusyWorker(extension, statuses)).toBe(statuses[1]);
|
||||
});
|
||||
|
||||
it("selectLeastBusyWorker returns only busy worker when queue length is 1", () => {
|
||||
const extension = {};
|
||||
const statuses = [{ worker: { id: "w1" }, queueRemaining: 3 }];
|
||||
expect(selectLeastBusyWorker(extension, statuses)).toBe(statuses[0]);
|
||||
});
|
||||
|
||||
it("buildWorkerWebSocketUrl converts http/https to ws/wss", () => {
|
||||
expect(buildWorkerWebSocketUrl("http://worker.local:8188")).toBe(
|
||||
"ws://worker.local:8188/distributed/worker_ws"
|
||||
|
||||
@@ -0,0 +1,242 @@
|
||||
import { afterEach, beforeEach, describe, expect, it } from "vitest";
|
||||
|
||||
import {
|
||||
buildWorkerUrl,
|
||||
buildWorkerWebSocketUrl,
|
||||
getMasterUrl,
|
||||
normalizeWorkerUrl,
|
||||
parseHostInput,
|
||||
} from "../urlUtils.js";
|
||||
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// normalizeWorkerUrl
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
describe("normalizeWorkerUrl", () => {
|
||||
it("returns empty string for empty input", () => {
|
||||
expect(normalizeWorkerUrl("")).toBe("");
|
||||
});
|
||||
|
||||
it("returns empty string for null input", () => {
|
||||
expect(normalizeWorkerUrl(null)).toBe("");
|
||||
});
|
||||
|
||||
it("returns empty string for non-string input", () => {
|
||||
expect(normalizeWorkerUrl(42)).toBe("");
|
||||
});
|
||||
|
||||
it("preserves https protocol", () => {
|
||||
expect(normalizeWorkerUrl("https://example.com")).toBe("https://example.com");
|
||||
});
|
||||
|
||||
it("strips trailing slash", () => {
|
||||
expect(normalizeWorkerUrl("http://example.com/")).toBe("http://example.com");
|
||||
});
|
||||
|
||||
it("prepends http when protocol is missing", () => {
|
||||
const result = normalizeWorkerUrl("worker.local:8188");
|
||||
expect(result).toMatch(/^http:\/\//);
|
||||
});
|
||||
|
||||
it("trims leading and trailing whitespace", () => {
|
||||
expect(normalizeWorkerUrl(" http://example.com ")).toBe("http://example.com");
|
||||
});
|
||||
});
|
||||
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// parseHostInput
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
describe("parseHostInput", () => {
|
||||
it("returns empty host and null port for null", () => {
|
||||
expect(parseHostInput(null)).toEqual({ host: "", port: null });
|
||||
});
|
||||
|
||||
it("returns empty host and null port for empty string", () => {
|
||||
expect(parseHostInput("")).toEqual({ host: "", port: null });
|
||||
});
|
||||
|
||||
it("parses hostname without port", () => {
|
||||
const result = parseHostInput("worker.example.com");
|
||||
expect(result.host).toBe("worker.example.com");
|
||||
expect(result.port).toBeNull();
|
||||
});
|
||||
|
||||
it("parses hostname with port", () => {
|
||||
const result = parseHostInput("worker.example.com:9000");
|
||||
expect(result.host).toBe("worker.example.com");
|
||||
expect(result.port).toBe(9000);
|
||||
});
|
||||
|
||||
it("strips http:// protocol prefix", () => {
|
||||
const result = parseHostInput("http://worker.example.com:8188");
|
||||
expect(result.host).toBe("worker.example.com");
|
||||
expect(result.port).toBe(8188);
|
||||
});
|
||||
|
||||
it("strips https:// protocol prefix", () => {
|
||||
const result = parseHostInput("https://worker.example.com");
|
||||
expect(result.host).toBe("worker.example.com");
|
||||
expect(result.port).toBeNull();
|
||||
});
|
||||
|
||||
it("ignores path after host:port", () => {
|
||||
const result = parseHostInput("worker.example.com:8188/some/path");
|
||||
expect(result.host).toBe("worker.example.com");
|
||||
expect(result.port).toBe(8188);
|
||||
});
|
||||
});
|
||||
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// buildWorkerWebSocketUrl
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
describe("buildWorkerWebSocketUrl", () => {
|
||||
it("converts http to ws", () => {
|
||||
expect(buildWorkerWebSocketUrl("http://worker.local:8188")).toBe(
|
||||
"ws://worker.local:8188/distributed/worker_ws"
|
||||
);
|
||||
});
|
||||
|
||||
it("converts https to wss", () => {
|
||||
expect(buildWorkerWebSocketUrl("https://worker.example.com")).toBe(
|
||||
"wss://worker.example.com/distributed/worker_ws"
|
||||
);
|
||||
});
|
||||
|
||||
it("always appends /distributed/worker_ws", () => {
|
||||
const url = buildWorkerWebSocketUrl("http://worker.local:8188");
|
||||
expect(url.endsWith("/distributed/worker_ws")).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// buildWorkerUrl (requires window.location stub)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
describe("buildWorkerUrl", () => {
|
||||
let originalWindow;
|
||||
|
||||
beforeEach(() => {
|
||||
originalWindow = globalThis.window;
|
||||
globalThis.window = {
|
||||
location: {
|
||||
hostname: "127.0.0.1",
|
||||
protocol: "http:",
|
||||
port: "8188",
|
||||
origin: "http://127.0.0.1:8188",
|
||||
},
|
||||
};
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
globalThis.window = originalWindow;
|
||||
});
|
||||
|
||||
it("builds local worker URL using window hostname when no host set", () => {
|
||||
const worker = { id: "w1", port: 8189 };
|
||||
expect(buildWorkerUrl(worker, "/prompt")).toBe("http://127.0.0.1:8189/prompt");
|
||||
});
|
||||
|
||||
it("builds remote worker URL using explicit host", () => {
|
||||
const worker = { id: "w2", host: "worker.example.com", port: 9000 };
|
||||
expect(buildWorkerUrl(worker, "/prompt")).toBe("http://worker.example.com:9000/prompt");
|
||||
});
|
||||
|
||||
it("builds cloud worker URL with https when type=cloud", () => {
|
||||
const worker = { id: "w3", host: "cloud.example.com", port: 443, type: "cloud" };
|
||||
expect(buildWorkerUrl(worker, "/prompt")).toBe("https://cloud.example.com/prompt");
|
||||
});
|
||||
|
||||
it("uses https for port 443 even without type=cloud", () => {
|
||||
const worker = { id: "w4", host: "worker.example.com", port: 443 };
|
||||
const result = buildWorkerUrl(worker, "");
|
||||
expect(result.startsWith("https://")).toBe(true);
|
||||
});
|
||||
|
||||
it("rewrites runpod proxy hostname for local port", () => {
|
||||
globalThis.window = {
|
||||
location: {
|
||||
hostname: "podabc.proxy.runpod.net",
|
||||
protocol: "https:",
|
||||
port: "",
|
||||
origin: "https://podabc.proxy.runpod.net",
|
||||
},
|
||||
};
|
||||
const worker = { id: "w5", port: 8189 };
|
||||
expect(buildWorkerUrl(worker, "/prompt")).toBe(
|
||||
"https://podabc-8189.proxy.runpod.net/prompt"
|
||||
);
|
||||
});
|
||||
|
||||
it("returns URL without trailing slash when no endpoint given", () => {
|
||||
const worker = { id: "w6", host: "worker.example.com", port: 8188 };
|
||||
const result = buildWorkerUrl(worker, "");
|
||||
expect(result.endsWith("/")).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// getMasterUrl
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
describe("getMasterUrl", () => {
|
||||
const _loc = (hostname, protocol = "http:", port = "8188") => ({
|
||||
hostname,
|
||||
protocol,
|
||||
port,
|
||||
origin: `${protocol}//${hostname}${port ? `:${port}` : ""}`,
|
||||
});
|
||||
|
||||
it("returns origin when master host not configured and hostname is non-localhost", () => {
|
||||
const loc = _loc("192.168.1.10");
|
||||
const result = getMasterUrl({}, loc);
|
||||
expect(result).toBe(loc.origin);
|
||||
});
|
||||
|
||||
it("returns origin for localhost when master host not configured", () => {
|
||||
const loc = _loc("127.0.0.1");
|
||||
const result = getMasterUrl({}, loc);
|
||||
expect(result).toBe(loc.origin);
|
||||
});
|
||||
|
||||
it("uses configured master host as-is when it includes http://", () => {
|
||||
const config = { master: { host: "http://master.example.com" } };
|
||||
const result = getMasterUrl(config, _loc("127.0.0.1"));
|
||||
expect(result).toBe("http://master.example.com");
|
||||
});
|
||||
|
||||
it("uses configured master host as-is when it includes https://", () => {
|
||||
const config = { master: { host: "https://secure.master.com" } };
|
||||
const result = getMasterUrl(config, _loc("127.0.0.1"));
|
||||
expect(result).toBe("https://secure.master.com");
|
||||
});
|
||||
|
||||
it("defaults to https for domain-name master hosts", () => {
|
||||
const config = { master: { host: "master.example.com" } };
|
||||
const result = getMasterUrl(config, _loc("127.0.0.1"));
|
||||
expect(result).toBe("https://master.example.com");
|
||||
});
|
||||
|
||||
it("does not force https for IP-address master hosts", () => {
|
||||
const config = { master: { host: "192.168.1.100" } };
|
||||
const result = getMasterUrl(config, _loc("127.0.0.1"));
|
||||
expect(result.startsWith("https://")).toBe(false);
|
||||
});
|
||||
|
||||
it("does not force https for localhost master host", () => {
|
||||
const config = { master: { host: "localhost" } };
|
||||
const result = getMasterUrl(config, _loc("127.0.0.1"));
|
||||
expect(result.startsWith("https://")).toBe(false);
|
||||
});
|
||||
|
||||
it("accepts null log parameter without throwing", () => {
|
||||
const loc = _loc("127.0.0.1");
|
||||
expect(() => getMasterUrl({}, loc, null)).not.toThrow();
|
||||
});
|
||||
});
|
||||
@@ -55,7 +55,6 @@ const cardConfigs = {
|
||||
dynamic: true,
|
||||
initialColor: (data) => data.enabled ? STATUS_COLORS.OFFLINE_RED : STATUS_COLORS.DISABLED_GRAY,
|
||||
initialTitle: (data) => data.enabled ? "Checking status..." : "Disabled",
|
||||
pulsing: (data) => data.enabled,
|
||||
id: (data) => `status-${data.id}`
|
||||
},
|
||||
infoText: (data, extension) => {
|
||||
|
||||
Reference in New Issue
Block a user