import base64 import importlib.util import io import sys import tempfile import types import unittest from pathlib import Path from unittest.mock import AsyncMock, patch from PIL import Image ROOT = Path(__file__).resolve().parents[1] MODULE_PATH = ROOT / "api" / "orchestration" / "serverless_media.py" def load_module(): spec = importlib.util.spec_from_file_location("dist_serverless_media_test", MODULE_PATH) assert spec is not None and spec.loader is not None module = importlib.util.module_from_spec(spec) sys.modules[spec.name] = module spec.loader.exec_module(module) return module media = load_module() def load_runpod_module(): spec = importlib.util.spec_from_file_location( "dist_serverless_media_runpod_test", ROOT / "api" / "orchestration" / "runpod_client.py" ) assert spec is not None and spec.loader is not None module = importlib.util.module_from_spec(spec) sys.modules[spec.name] = module spec.loader.exec_module(module) return module runpod = load_runpod_module() def png_bytes(mode="RGB", size=(2, 1), color=None): if color is None: color = (10, 20, 30, 40) if mode == "RGBA" else (10, 20, 30) buffer = io.BytesIO() Image.new(mode, size, color).save(buffer, format="PNG") return buffer.getvalue() class FolderPathsFixture: def __init__(self, root): self.input = root / "input" self.output = root / "output" self.temp = root / "temp" for directory in (self.input, self.output, self.temp): directory.mkdir(parents=True, exist_ok=True) def get_input_directory(self): return str(self.input) def get_output_directory(self): return str(self.output) def get_temp_directory(self): return str(self.temp) def get_annotated_filepath(self, filename): value = filename.replace("\\", "/") # Match real ComfyUI folder_paths: root annotations are suffixes, # not directory prefixes (output/foo.png is an input subdirectory). for marker, root in (("input", self.input), ("output", self.output), ("temp", self.temp)): suffix = f" [{marker}]" if value.endswith(suffix): return str(root / value[:-len(suffix)]) return str(self.input / value) class ServerlessMediaInputTests(unittest.IsolatedAsyncioTestCase): async def asyncSetUp(self): self.tempdir = tempfile.TemporaryDirectory() self.paths = FolderPathsFixture(Path(self.tempdir.name)) self.folder_paths = types.ModuleType("folder_paths") self.folder_paths.get_input_directory = self.paths.get_input_directory self.folder_paths.get_output_directory = self.paths.get_output_directory self.folder_paths.get_temp_directory = self.paths.get_temp_directory self.folder_paths.get_annotated_filepath = self.paths.get_annotated_filepath self.folder_patch = patch.dict(sys.modules, {"folder_paths": self.folder_paths}) self.folder_patch.start() self.addCleanup(self.folder_patch.stop) self.addCleanup(self.tempdir.cleanup) async def test_packs_image_and_rewrites_duplicate_loadimage_inputs(self): payload = png_bytes() (self.paths.input / "cat.png").write_bytes(payload) workflow = { "1": {"class_type": "LoadImage", "inputs": {"image": "cat.png [input]", "other": "unchanged"}}, "2": {"class_type": "LoadImage", "inputs": {"image": "cat.png [input]"}}, "3": {"class_type": "SaveImage", "inputs": {"filename_prefix": "stable"}}, } packed = await media.prepare_images(workflow) self.assertEqual(len(packed), 1) self.assertEqual(set(packed[0]), {"name", "image"}) self.assertRegex(packed[0]["name"], r"^[A-Za-z0-9_.-]+\.png$") self.assertEqual(base64.b64decode(packed[0]["image"], validate=True), payload) self.assertEqual(workflow["1"]["inputs"]["image"], packed[0]["name"]) self.assertEqual(workflow["2"]["inputs"]["image"], packed[0]["name"]) self.assertEqual(workflow["1"]["inputs"]["other"], "unchanged") self.assertEqual(workflow["3"]["inputs"]["filename_prefix"], "stable") async def test_resolves_annotated_input_output_and_temp_sources(self): sources = ((self.paths.input, "one.png"), (self.paths.output, "two.jpg"), (self.paths.temp, "three.webp")) workflow = {} for index, (directory, filename) in enumerate(sources, 1): (directory / filename).write_bytes(png_bytes()) root = directory.name workflow[str(index)] = {"class_type": "LoadImage", "inputs": {"image": f"{filename} [{root}]"}} packed = await media.prepare_images(workflow) self.assertEqual(len(packed), 3) self.assertEqual({base64.b64decode(item["image"]) for item in packed}, {png_bytes()}) self.assertEqual({node["inputs"]["image"] for node in workflow.values()}, {item["name"] for item in packed}) async def test_directory_prefix_is_not_a_root_annotation(self): (self.paths.input / "output").mkdir() payload = png_bytes(color=(100, 0, 0)) (self.paths.input / "output" / "same.png").write_bytes(payload) (self.paths.output / "same.png").write_bytes(png_bytes()) workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": "output/same.png"}}} packed = await media.prepare_images(workflow) self.assertEqual(base64.b64decode(packed[0]["image"]), payload) async def test_rejects_unknown_root_annotation(self): (self.paths.input / "cat.png").write_bytes(png_bytes()) workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": "cat.png [abc123]"}}} with self.assertRaisesRegex(media.ServerlessMediaError, "annotation"): await media.prepare_images(workflow) async def test_recognizes_custom_file_image_reference(self): (self.paths.input / "source.bmp").write_bytes(png_bytes()) workflow = {"1": {"class_type": "CustomImageLoader", "inputs": {"file": "source.bmp"}}} packed = await media.prepare_images(workflow) self.assertEqual(len(packed), 1) self.assertEqual(workflow["1"]["inputs"]["file"], packed[0]["name"]) async def test_rejects_missing_and_dynamic_loadimage_values(self): cases = ( {"1": {"class_type": "LoadImage", "inputs": {}}}, {"1": {"class_type": "LoadImage"}}, {"1": {"class_type": "LoadImage", "inputs": None}}, {"1": {"class_type": "LoadImage", "inputs": {"image": ["2", 0]}}}, ) for workflow in cases: with self.subTest(workflow=workflow): with self.assertRaisesRegex(media.ServerlessMediaError, "LoadImage"): await media.prepare_images(workflow) async def test_rejects_missing_source_and_unapproved_extension(self): for value in ("absent.png", "animation.gif", "clip.mp4"): with self.subTest(value=value): workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": value}}} with self.assertRaises(media.ServerlessMediaError): await media.prepare_images(workflow) async def test_rejects_unsupported_audio_or_video_references_without_reading(self): for key, value in (("audio", "private.wav"), ("video", "private.mp4"), ("file", "private.wav")): with self.subTest(key=key): workflow = {"1": {"class_type": "CustomLoader", "inputs": {key: value}}} with self.assertRaisesRegex(media.ServerlessMediaError, "audio|video|media"): await media.prepare_images(workflow) async def test_rejects_parent_traversal_and_absolute_outside_roots(self): outside = Path(self.tempdir.name).parent / "outside.png" outside.write_bytes(png_bytes()) self.addCleanup(lambda: outside.unlink(missing_ok=True)) for value in ("../outside.png", str(outside)): with self.subTest(value=value): workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": value}}} with self.assertRaises(media.ServerlessMediaError): await media.prepare_images(workflow) async def test_rejects_symlink_escape_from_allowed_root(self): outside = Path(self.tempdir.name).parent / "linked-outside.png" outside.write_bytes(png_bytes()) self.addCleanup(lambda: outside.unlink(missing_ok=True)) (self.paths.input / "link.png").symlink_to(outside) workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": "link.png"}}} with self.assertRaises(media.ServerlessMediaError): await media.prepare_images(workflow) async def test_rejects_individual_and_aggregate_input_size_overages(self): huge = self.paths.input / "large.jpg" with huge.open("wb") as handle: handle.truncate(media.MAX_INPUT_IMAGE_BYTES + 1) workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": "large.jpg"}}} with self.assertRaisesRegex(media.ServerlessMediaError, "size"): await media.prepare_images(workflow) huge.unlink() first = self.paths.input / "first.png" second = self.paths.input / "second.png" first.write_bytes(b"a" * (media.MAX_INPUT_TOTAL_BYTES // 2 + 1)) second.write_bytes(b"b" * (media.MAX_INPUT_TOTAL_BYTES // 2 + 1)) workflow = { "1": {"class_type": "LoadImage", "inputs": {"image": "first.png"}}, "2": {"class_type": "LoadImage", "inputs": {"image": "second.png"}}, } with self.assertRaisesRegex(media.ServerlessMediaError, "size"): await media.prepare_images(workflow) async def test_accepts_raw_image_over_old_7mib_when_complete_request_fits_10mb(self): raw_size = 7 * 1024 * 1024 + 1 source = self.paths.input / "large-input.png" source.write_bytes(b"x" * raw_size) workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": "large-input.png"}}} packed = await media.prepare_images(workflow) request_body = runpod.serialize_request(workflow, packed, 30) self.assertGreater(raw_size, 7 * 1024 * 1024) self.assertEqual(len(base64.b64decode(packed[0]["image"])), raw_size) self.assertLessEqual(len(request_body), 10_000_000) class ServerlessMediaOutputTests(unittest.IsolatedAsyncioTestCase): def manifest(self, prefix="render"): return [{"collector_id": "collector-1", "job_id": "job-1", "node_id": "12", "kind": "image", "prefix": prefix}] def artifact(self, filename, payload=None, **extra): entry = { "filename": filename, "artifact_type": "output", "artifact_path": f"/safe/output/{filename}", "subfolder": "", "type": "base64", "data": base64.b64encode(payload if payload is not None else png_bytes()).decode("ascii"), } entry.update(extra) return entry async def test_decodes_rgb_to_normalized_float32_and_orders_by_counter(self): output = {"images": [self.artifact("render_00009_.png"), self.artifact("render_00002_.png")]} result = await media.collect_images(output, self.manifest()) self.assertEqual(list(result), ["job-1"]) self.assertEqual(len(result["job-1"]), 2) self.assertTrue(all(image.dtype.name == "float32" for image in result["job-1"])) self.assertEqual(result["job-1"][0].shape, (1, 2, 3)) self.assertAlmostEqual(float(result["job-1"][0][0, 0, 0]), 10 / 255) async def test_groups_images_under_manifest_job_ids(self): manifest = [ self.manifest("first")[0], {"collector_id": "collector-2", "job_id": "job-2", "node_id": "14", "kind": "image", "prefix": "second"}, ] output = {"images": [self.artifact("first_00001_.png"), self.artifact("second_00001_.png")]} result = await media.collect_images(output, manifest) self.assertEqual(set(result), {"job-1", "job-2"}) self.assertEqual(len(result["job-1"]), 1) self.assertEqual(len(result["job-2"]), 1) async def test_preserves_rgba_alpha(self): output = {"images": [self.artifact("render_00001_.png", png_bytes(mode="RGBA"))]} result = await media.collect_images(output, self.manifest()) self.assertEqual(result["job-1"][0].shape, (1, 2, 4)) self.assertAlmostEqual(float(result["job-1"][0][0, 0, 3]), 40 / 255) async def test_signed_url_is_downloaded_only_through_injected_downloader(self): output = {"images": [{ "filename": "render_00001_.png", "artifact_type": "output", "type": "gcs_url", "url": "https://storage.example.test/signed", }]} fetch = AsyncMock(return_value=png_bytes()) result = await media.collect_images(output, self.manifest(), download=fetch) self.assertEqual(result["job-1"][0].shape, (1, 2, 3)) fetch.assert_awaited_once_with("https://storage.example.test/signed") async def test_ignores_unrelated_artifacts_but_requires_every_manifest_collector(self): output = {"images": [ self.artifact("preview_00001_.png"), self.artifact("render_00001_.png"), {"filename": "mesh.obj", "type": "url", "url": "not a fetch target"}, ]} result = await media.collect_images(output, self.manifest()) self.assertEqual(len(result["job-1"]), 1) with self.assertRaisesRegex(media.ServerlessMediaError, "expected|missing"): await media.collect_images({"images": [self.artifact("preview_00001_.png")]}, self.manifest()) async def test_accepts_compiler_length_flat_safe_prefix(self): prefix = "serverless_" + "j" * 24 + "_" + "w" * 24 + "_" + "c" * 24 + "_image_" + "a" * 64 manifest = self.manifest(prefix) output = {"images": [self.artifact(f"{prefix}_00001_.png")]} result = await media.collect_images(output, manifest) self.assertEqual(len(result["job-1"]), 1) async def test_rejects_duplicate_manifest_prefix_and_unsafe_prefix(self): duplicate = self.manifest() * 2 with self.assertRaises(media.ServerlessMediaError): await media.collect_images({"images": []}, duplicate) unsafe = self.manifest("../render") with self.assertRaises(media.ServerlessMediaError): await media.collect_images({"images": []}, unsafe) async def test_rejects_duplicate_canonical_artifact_identity(self): item = self.artifact("render_00001_.png") with self.assertRaisesRegex(media.ServerlessMediaError, "[Dd]uplicate"): await media.collect_images({"images": [item, dict(item)]}, self.manifest()) async def test_rejects_matching_non_output_artifact_and_non_numeric_counter(self): wrong_type = self.artifact("render_00001_.png", artifact_type="temp") with self.assertRaises(media.ServerlessMediaError): await media.collect_images({"images": [wrong_type]}, self.manifest()) with self.assertRaises(media.ServerlessMediaError): await media.collect_images({"images": [self.artifact("render_final.png")]}, self.manifest()) async def test_rejects_provider_errors_even_when_job_completed(self): for output in ( {"error": "raw provider secret", "images": []}, {"errors": ["raw provider secret"], "images": []}, {"status": "COMPLETED", "output": {"error": "raw provider secret", "images": []}}, ): with self.subTest(output=output): with self.assertRaises(media.ServerlessMediaError) as ctx: await media.collect_images(output, self.manifest()) self.assertNotIn("raw provider secret", str(ctx.exception)) async def test_rejects_invalid_base64_and_malformed_image(self): for item in ( self.artifact("render_00001_.png", data="!not base64!"), self.artifact("render_00001_.png", data=base64.b64encode(b"not an image").decode()), ): with self.subTest(item=item): with self.assertRaises(media.ServerlessMediaError): await media.collect_images({"images": [item]}, self.manifest()) async def test_validates_shape_consistency_within_collector(self): output = {"images": [ self.artifact("render_00001_.png", png_bytes(size=(2, 1))), self.artifact("render_00002_.png", png_bytes(size=(3, 1))), ]} with self.assertRaisesRegex(media.ServerlessMediaError, "shape|dimensions"): await media.collect_images(output, self.manifest()) async def test_bounds_output_file_count_and_bytes(self): too_many = {"images": [{}] * (media.MAX_OUTPUT_FILES + 1)} with self.assertRaisesRegex(media.ServerlessMediaError, "files"): await media.collect_images(too_many, self.manifest()) with patch.object(media, "MAX_OUTPUT_IMAGE_BYTES", 4): output = {"images": [self.artifact("render_00001_.png", data=base64.b64encode(b"12345").decode())]} with self.assertRaisesRegex(media.ServerlessMediaError, "size"): await media.collect_images(output, self.manifest()) async def test_aggregate_payload_limit_is_checked_as_artifacts_are_read(self): output = {"images": [ self.artifact("render_00001_.png", data=base64.b64encode(b"12").decode()), self.artifact("render_00002_.png", data=base64.b64encode(b"34").decode()), ]} with patch.object(media, "MAX_OUTPUT_TOTAL_BYTES", 3): with self.assertRaisesRegex(media.ServerlessMediaError, "size"): await media.collect_images(output, self.manifest()) async def test_rejects_non_public_or_non_https_artifact_urls_before_downloading(self): for url in ("http://example.test/a", "https://127.0.0.1/a", "https://user:pass@example.test/a"): with self.subTest(url=url): fetch = AsyncMock(return_value=png_bytes()) output = {"images": [{"filename": "render_00001_.png", "type": "url", "url": url}]} with self.assertRaises(media.ServerlessMediaError): await media.collect_images(output, self.manifest(), download=fetch) fetch.assert_not_awaited() async def test_public_resolver_rejects_non_global_dns_answers(self): resolver = media._PublicResolver() resolver._getaddrinfo = AsyncMock(return_value=[(2, 1, 6, "", ("8.8.8.8", 443))]) self.assertEqual(len(await resolver.resolve("cdn.example.test", 443)), 1) resolver._getaddrinfo = AsyncMock(return_value=[(2, 1, 6, "", ("10.0.0.5", 443))]) with self.assertRaises(OSError): await resolver.resolve("cdn.example.test", 443) async def test_default_download_revalidates_redirects_and_disables_environment_proxies(self): class Content: def __init__(self, chunks): self.chunks = iter(chunks) async def iter_chunked(self, _size): for chunk in self.chunks: yield chunk class Response: def __init__(self, status, headers=None, chunks=()): self.status, self.headers = status, headers or {} self.content = Content(chunks) async def __aenter__(self): return self async def __aexit__(self, *_args): return False class Session: def __init__(self, responses): self.responses, self.calls = list(responses), [] async def __aenter__(self): return self async def __aexit__(self, *_args): return False def get(self, url, **kwargs): self.calls.append((url, kwargs)) return self.responses.pop(0) session = Session([Response(302, {"Location": "https://cdn.example.test/final"}), Response(200, chunks=[b"abc", b"def"])]) with patch.object(media.aiohttp, "ClientSession", return_value=session) as create_session: payload = await media.download_public_artifact("https://storage.example.test/signed") self.assertEqual(payload, b"abcdef") self.assertEqual([call[0] for call in session.calls], ["https://storage.example.test/signed", "https://cdn.example.test/final"]) self.assertTrue(all(call[1]["allow_redirects"] is False for call in session.calls)) self.assertFalse(create_session.call_args.kwargs["trust_env"]) self.assertIsInstance(create_session.call_args.kwargs["connector"]._resolver, media._PublicResolver) self.assertNotIn("Authorization", session.calls[0][1].get("headers", {})) await create_session.call_args.kwargs["connector"].close() blocked = Session([Response(302, {"Location": "https://127.0.0.1/metadata"})]) with patch.object(media.aiohttp, "ClientSession", return_value=blocked) as blocked_session: with self.assertRaises(media.ServerlessMediaError): await media.download_public_artifact("https://storage.example.test/signed") self.assertEqual(len(blocked.calls), 1) await blocked_session.call_args.kwargs["connector"].close() if __name__ == "__main__": unittest.main()