448 lines
21 KiB
Python
448 lines
21 KiB
Python
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()
|