From 1630a1ab791dab5444ebfba372ea06e69d5b38ec Mon Sep 17 00:00:00 2001 From: Peyton DeNiro <25550995+pydn@users.noreply.github.com> Date: Sun, 29 Mar 2026 14:08:18 -0500 Subject: [PATCH] Strengthen test suite robustness --- tests/README.md | 16 +- .../runtime/reused-node-class-branches.json | 27 ++ .../runtime/secondary-output-selection.json | 17 + tests/runtime/generated/text-to-image.py | 235 ++++++++++++++ .../runtime/generated/upscale-model-loader.py | 200 ++++++++++++ tests/runtime/run_runtime_validation.py | 50 +++ tests/test_runtime_validation_harness.py | 305 ++++++++++++++++++ tests/test_upscale_model_loader_export.py | 112 ++++++- 8 files changed, 957 insertions(+), 5 deletions(-) create mode 100644 tests/fixtures/runtime/reused-node-class-branches.json create mode 100644 tests/fixtures/runtime/secondary-output-selection.json create mode 100644 tests/runtime/generated/text-to-image.py create mode 100644 tests/runtime/generated/upscale-model-loader.py create mode 100644 tests/test_runtime_validation_harness.py diff --git a/tests/README.md b/tests/README.md index cc7bca6..af5dfd2 100644 --- a/tests/README.md +++ b/tests/README.md @@ -5,6 +5,8 @@ This directory contains two test paths: - unit tests for exporter behavior - runtime validation for committed workflow fixtures +Use the lightweight path for routine contributor validation, and use the runtime path when you need end-to-end confidence against a real ComfyUI checkout. + ## Prerequisites - Run commands from the repo root. @@ -30,12 +32,18 @@ The current runtime-capable fixtures require these model files in the target Com ## Unit Tests -Run the current unit test module: +Run the exporter-focused unit test module: ```bash uv run python -m unittest tests.test_upscale_model_loader_export ``` +Run the runtime-harness unit test module: + +```bash +uv run python -m unittest tests.test_runtime_validation_harness +``` + Run all `unittest`-discoverable tests under `tests`: ```bash @@ -49,6 +57,7 @@ The runtime harness lives at `tests/runtime/run_runtime_validation.py`. ### Fast Tier Fast tier validates export behavior against committed fixtures without requiring a full ComfyUI runtime for every fixture. +This is the recommended default validation lane for routine changes. Run all fast-tier compatible fixtures: @@ -65,6 +74,7 @@ uv run python tests/runtime/run_runtime_validation.py --tier fast --fixture unsa ### Runtime Tier Runtime tier exports inside a real ComfyUI checkout and executes generated Python for runtime-capable fixtures. +This is the heavier validation lane for changes that need end-to-end runtime confidence. Run all runtime-capable fixtures: @@ -92,11 +102,15 @@ Current committed fixtures: - `text-to-image` - `unsafe-kwargs` - `subgraph-identifiers` +- `reused-node-class-branches` +- `secondary-output-selection` Notes: - `--tier runtime` only runs fixtures marked runtime-capable. - `--tier fast` only runs fixtures with local test mappings. +- `reused-node-class-branches` protects repeated node-class usage and branch wiring in the fast tier. +- `secondary-output-selection` protects non-zero output index wiring in the fast tier. ## Troubleshooting diff --git a/tests/fixtures/runtime/reused-node-class-branches.json b/tests/fixtures/runtime/reused-node-class-branches.json new file mode 100644 index 0000000..27c61ec --- /dev/null +++ b/tests/fixtures/runtime/reused-node-class-branches.json @@ -0,0 +1,27 @@ +{ + "1:10": { + "class_type": "PassthroughText", + "inputs": { + "text": "left branch" + } + }, + "1:20": { + "class_type": "PassthroughText", + "inputs": { + "text": "right branch" + } + }, + "2": { + "class_type": "JoinText", + "inputs": { + "left": [ + "1:10", + 0 + ], + "right": [ + "1:20", + 0 + ] + } + } +} diff --git a/tests/fixtures/runtime/secondary-output-selection.json b/tests/fixtures/runtime/secondary-output-selection.json new file mode 100644 index 0000000..db9a055 --- /dev/null +++ b/tests/fixtures/runtime/secondary-output-selection.json @@ -0,0 +1,17 @@ +{ + "1": { + "class_type": "SplitText", + "inputs": { + "text": "alpha|omega" + } + }, + "2": { + "class_type": "PassthroughText", + "inputs": { + "text": [ + "1", + 1 + ] + } + } +} diff --git a/tests/runtime/generated/text-to-image.py b/tests/runtime/generated/text-to-image.py new file mode 100644 index 0000000..68568a2 --- /dev/null +++ b/tests/runtime/generated/text-to-image.py @@ -0,0 +1,235 @@ +# Imports +import json +import os +import random +import sys +from typing import Sequence, Mapping, Any, Union + + +def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any: + """Return a sequence or mapping result item by index.""" + try: + return obj[index] + except KeyError: + return obj["result"][index] + + +def get_comfyui_path() -> str: + """Return the configured ComfyUI path, preferring COMFYUI_PATH when set.""" + comfyui_path = os.environ.get("COMFYUI_PATH") + if comfyui_path: + return comfyui_path + return find_path("ComfyUI") + + +def find_path(name: str, path: str = None) -> str: + """Recursively search parent folders until the named entry is found.""" + if path is None: + path = os.getcwd() + + if name in os.listdir(path): + path_name = os.path.join(path, name) + print(f"{name} found: {path_name}") + return path_name + + parent_directory = os.path.dirname(path) + if parent_directory == path: + return None + + return find_path(name, parent_directory) + + +def add_comfyui_directory_to_sys_path() -> None: + """Add the ComfyUI checkout to sys.path.""" + comfyui_path = get_comfyui_path() + if comfyui_path is not None and os.path.isdir(comfyui_path): + if comfyui_path in sys.path: + sys.path.remove(comfyui_path) + sys.path.insert(0, comfyui_path) + print(f"'{comfyui_path}' added to sys.path") + + +def add_extra_model_paths() -> None: + """Load ComfyUI extra model paths configuration when available.""" + try: + from main import load_extra_path_config + except ImportError: + print( + "Could not import load_extra_path_config from main.py. Looking in utils.extra_config instead." + ) + from utils.extra_config import load_extra_path_config + + extra_model_paths = find_path("extra_model_paths.yaml") + if extra_model_paths is not None: + load_extra_path_config(extra_model_paths) + else: + print("Could not find the extra_model_paths config file.") + + +def bootstrap_comfyui_runtime() -> None: + """Mirror the allocator-related ComfyUI startup steps before torch import.""" + add_comfyui_directory_to_sys_path() + + import comfy.options + + comfy.options.enable_args_parsing() + + from comfy.cli_args import args + + if os.name == "nt": + os.environ["MIMALLOC_PURGE_DELAY"] = "0" + + if args.default_device is not None: + default_dev = args.default_device + devices = list(range(32)) + devices.remove(default_dev) + devices.insert(0, default_dev) + devices = ",".join(map(str, devices)) + os.environ["CUDA_VISIBLE_DEVICES"] = str(devices) + os.environ["HIP_VISIBLE_DEVICES"] = str(devices) + + if args.cuda_device is not None: + os.environ["CUDA_VISIBLE_DEVICES"] = str(args.cuda_device) + os.environ["HIP_VISIBLE_DEVICES"] = str(args.cuda_device) + os.environ["ASCEND_RT_VISIBLE_DEVICES"] = str(args.cuda_device) + + if args.oneapi_device_selector is not None: + os.environ["ONEAPI_DEVICE_SELECTOR"] = args.oneapi_device_selector + + if args.deterministic and "CUBLAS_WORKSPACE_CONFIG" not in os.environ: + os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" + + import cuda_malloc + + if "rocm" in cuda_malloc.get_torch_version_noimport(): + os.environ["OCL_SET_SVM_SIZE"] = "262144" + + +# Runtime support + +bootstrap_comfyui_runtime() +import torch + +add_extra_model_paths() + + +# Workflow data +def build_workflow() -> dict[str, Any]: + return { + "1": { + "class_type": "CheckpointLoaderSimple", + "inputs": {"ckpt_name": "v1-5-pruned-emaonly-fp16.safetensors"}, + }, + "2": { + "class_type": "CLIPTextEncode", + "inputs": { + "text": "a small cottage in a meadow, soft daylight", + "clip": ["1", 1], + }, + }, + "3": { + "class_type": "CLIPTextEncode", + "inputs": {"text": "blurry, low quality", "clip": ["1", 1]}, + }, + "4": { + "class_type": "EmptyLatentImage", + "inputs": {"width": 512, "height": 512, "batch_size": 1}, + }, + "5": { + "class_type": "KSampler", + "inputs": { + "seed": 1, + "steps": 4, + "cfg": 7, + "sampler_name": "euler", + "scheduler": "normal", + "denoise": 1, + "model": ["1", 0], + "positive": ["2", 0], + "negative": ["3", 0], + "latent_image": ["4", 0], + }, + }, + "6": { + "class_type": "VAEDecode", + "inputs": {"samples": ["5", 0], "vae": ["1", 2]}, + }, + "7": { + "class_type": "SaveImage", + "inputs": {"filename_prefix": "E2E_text_to_image", "images": ["6", 0]}, + }, + } + + +def build_extra_pnginfo() -> dict[str, Any] | None: + return None + + +workflow = build_workflow() +prompt = json.loads(json.dumps(workflow)) +extra_pnginfo = build_extra_pnginfo() + +# Workflow execution +# Node imports +from nodes import ( + CLIPTextEncode, + CheckpointLoaderSimple, + EmptyLatentImage, + KSampler, + NODE_CLASS_MAPPINGS, + SaveImage, + VAEDecode, +) + + +def main(): + with torch.inference_mode(): + checkpointloadersimple = CheckpointLoaderSimple() + checkpointloadersimple_1 = checkpointloadersimple.load_checkpoint( + ckpt_name="v1-5-pruned-emaonly-fp16.safetensors" + ) + cliptextencode = CLIPTextEncode() + cliptextencode_2 = cliptextencode.encode( + text="a small cottage in a meadow, soft daylight", + clip=get_value_at_index(checkpointloadersimple_1, 1), + ) + cliptextencode_3 = cliptextencode.encode( + text="blurry, low quality", + clip=get_value_at_index(checkpointloadersimple_1, 1), + ) + emptylatentimage = EmptyLatentImage() + emptylatentimage_4 = emptylatentimage.generate( + width=512, height=512, batch_size=1 + ) + ksampler = KSampler() + vaedecode = VAEDecode() + saveimage = SaveImage() + for q in range(1): + node_5_seed = prompt["5"]["inputs"]["seed"] = random.randint(1, 2**64) + ksampler_5 = ksampler.sample( + seed=node_5_seed, + steps=4, + cfg=7, + sampler_name="euler", + scheduler="normal", + denoise=1, + model=get_value_at_index(checkpointloadersimple_1, 0), + positive=get_value_at_index(cliptextencode_2, 0), + negative=get_value_at_index(cliptextencode_3, 0), + latent_image=get_value_at_index(emptylatentimage_4, 0), + ) + vaedecode_6 = vaedecode.decode( + samples=get_value_at_index(ksampler_5, 0), + vae=get_value_at_index(checkpointloadersimple_1, 2), + ) + saveimage_7 = saveimage.save_images( + filename_prefix="E2E_text_to_image", + images=get_value_at_index(vaedecode_6, 0), + prompt=prompt, + extra_pnginfo=extra_pnginfo, + ) + + +# Entrypoint +if __name__ == "__main__": + main() diff --git a/tests/runtime/generated/upscale-model-loader.py b/tests/runtime/generated/upscale-model-loader.py new file mode 100644 index 0000000..b7b0a09 --- /dev/null +++ b/tests/runtime/generated/upscale-model-loader.py @@ -0,0 +1,200 @@ +# Imports +import json +import os +import random +import sys +from typing import Sequence, Mapping, Any, Union + + +def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any: + """Return a sequence or mapping result item by index.""" + try: + return obj[index] + except KeyError: + return obj["result"][index] + + +def get_comfyui_path() -> str: + """Return the configured ComfyUI path, preferring COMFYUI_PATH when set.""" + comfyui_path = os.environ.get("COMFYUI_PATH") + if comfyui_path: + return comfyui_path + return find_path("ComfyUI") + + +def find_path(name: str, path: str = None) -> str: + """Recursively search parent folders until the named entry is found.""" + if path is None: + path = os.getcwd() + + if name in os.listdir(path): + path_name = os.path.join(path, name) + print(f"{name} found: {path_name}") + return path_name + + parent_directory = os.path.dirname(path) + if parent_directory == path: + return None + + return find_path(name, parent_directory) + + +def add_comfyui_directory_to_sys_path() -> None: + """Add the ComfyUI checkout to sys.path.""" + comfyui_path = get_comfyui_path() + if comfyui_path is not None and os.path.isdir(comfyui_path): + if comfyui_path in sys.path: + sys.path.remove(comfyui_path) + sys.path.insert(0, comfyui_path) + print(f"'{comfyui_path}' added to sys.path") + + +def add_extra_model_paths() -> None: + """Load ComfyUI extra model paths configuration when available.""" + try: + from main import load_extra_path_config + except ImportError: + print( + "Could not import load_extra_path_config from main.py. Looking in utils.extra_config instead." + ) + from utils.extra_config import load_extra_path_config + + extra_model_paths = find_path("extra_model_paths.yaml") + if extra_model_paths is not None: + load_extra_path_config(extra_model_paths) + else: + print("Could not find the extra_model_paths config file.") + + +def bootstrap_comfyui_runtime() -> None: + """Mirror the allocator-related ComfyUI startup steps before torch import.""" + add_comfyui_directory_to_sys_path() + + import comfy.options + + comfy.options.enable_args_parsing() + + from comfy.cli_args import args + + if os.name == "nt": + os.environ["MIMALLOC_PURGE_DELAY"] = "0" + + if args.default_device is not None: + default_dev = args.default_device + devices = list(range(32)) + devices.remove(default_dev) + devices.insert(0, default_dev) + devices = ",".join(map(str, devices)) + os.environ["CUDA_VISIBLE_DEVICES"] = str(devices) + os.environ["HIP_VISIBLE_DEVICES"] = str(devices) + + if args.cuda_device is not None: + os.environ["CUDA_VISIBLE_DEVICES"] = str(args.cuda_device) + os.environ["HIP_VISIBLE_DEVICES"] = str(args.cuda_device) + os.environ["ASCEND_RT_VISIBLE_DEVICES"] = str(args.cuda_device) + + if args.oneapi_device_selector is not None: + os.environ["ONEAPI_DEVICE_SELECTOR"] = args.oneapi_device_selector + + if args.deterministic and "CUBLAS_WORKSPACE_CONFIG" not in os.environ: + os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" + + import cuda_malloc + + if "rocm" in cuda_malloc.get_torch_version_noimport(): + os.environ["OCL_SET_SVM_SIZE"] = "262144" + + +# Runtime support + +bootstrap_comfyui_runtime() +import torch + +add_extra_model_paths() + + +def import_custom_nodes() -> None: + """Initialize ComfyUI custom nodes in the exporter runtime.""" + comfyui_path = get_comfyui_path() + if comfyui_path and comfyui_path not in sys.path: + sys.path.insert(0, comfyui_path) + + import asyncio + import execution + from nodes import init_extra_nodes + + if comfyui_path in sys.path: + sys.path.remove(comfyui_path) + sys.path.insert(0, comfyui_path) + + import server + + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + server_instance = server.PromptServer(loop) + execution.PromptQueue(server_instance) + asyncio.run(init_extra_nodes()) + + +# Workflow data +def build_workflow() -> dict[str, Any]: + return { + "1": {"class_type": "LoadImage", "inputs": {"image": "e2e_upscale_input.png"}}, + "2": { + "class_type": "UpscaleModelLoader", + "inputs": {"model_name": "RealESRGAN_x4plus.safetensors"}, + }, + "3": { + "class_type": "ImageUpscaleWithModel", + "inputs": {"upscale_model": ["2", 0], "image": ["1", 0]}, + }, + "4": { + "class_type": "SaveImage", + "inputs": { + "filename_prefix": "E2E_upscale_model_loader", + "images": ["3", 0], + }, + }, + } + + +def build_extra_pnginfo() -> dict[str, Any] | None: + return None + + +workflow = build_workflow() +prompt = json.loads(json.dumps(workflow)) +extra_pnginfo = build_extra_pnginfo() + +# Workflow execution +# Node imports +from nodes import LoadImage, NODE_CLASS_MAPPINGS, SaveImage + + +def main(): + with torch.inference_mode(): + import_custom_nodes() + loadimage = LoadImage() + loadimage_1 = loadimage.load_image(image="e2e_upscale_input.png") + upscalemodelloader = NODE_CLASS_MAPPINGS["UpscaleModelLoader"]() + upscalemodelloader_2 = upscalemodelloader.EXECUTE_NORMALIZED( + model_name="RealESRGAN_x4plus.safetensors" + ) + imageupscalewithmodel = NODE_CLASS_MAPPINGS["ImageUpscaleWithModel"]() + saveimage = SaveImage() + for q in range(1): + imageupscalewithmodel_3 = imageupscalewithmodel.EXECUTE_NORMALIZED( + upscale_model=get_value_at_index(upscalemodelloader_2, 0), + image=get_value_at_index(loadimage_1, 0), + ) + saveimage_4 = saveimage.save_images( + filename_prefix="E2E_upscale_model_loader", + images=get_value_at_index(imageupscalewithmodel_3, 0), + prompt=prompt, + extra_pnginfo=extra_pnginfo, + ) + + +# Entrypoint +if __name__ == "__main__": + main() diff --git a/tests/runtime/run_runtime_validation.py b/tests/runtime/run_runtime_validation.py index 46b0d84..235f99a 100644 --- a/tests/runtime/run_runtime_validation.py +++ b/tests/runtime/run_runtime_validation.py @@ -109,6 +109,40 @@ class PassthroughText: return (text,) +class SplitText: + CATEGORY = "utils" + FUNCTION = "split" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "text": ("STRING",), + } + } + + def split(self, text): + left, right = text.split("|", 1) + return (left, right) + + +class JoinText: + CATEGORY = "utils" + FUNCTION = "join" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "left": ("STRING",), + "right": ("STRING",), + } + } + + def join(self, left, right): + return (f"{left}::{right}",) + + class UpscaleModelLoader: CATEGORY = "loaders" FUNCTION = "load_model" @@ -234,6 +268,22 @@ FIXTURES = { "PassthroughText": PassthroughText, }, ), + "reused-node-class-branches": FixtureConfig( + name="reused-node-class-branches", + path=FIXTURE_DIR / "reused-node-class-branches.json", + mapping_factory=lambda: { + "PassthroughText": PassthroughText, + "JoinText": JoinText, + }, + ), + "secondary-output-selection": FixtureConfig( + name="secondary-output-selection", + path=FIXTURE_DIR / "secondary-output-selection.json", + mapping_factory=lambda: { + "SplitText": SplitText, + "PassthroughText": PassthroughText, + }, + ), } diff --git a/tests/test_runtime_validation_harness.py b/tests/test_runtime_validation_harness.py new file mode 100644 index 0000000..0e30253 --- /dev/null +++ b/tests/test_runtime_validation_harness.py @@ -0,0 +1,305 @@ +import struct +import tempfile +import unittest +import zlib +from pathlib import Path +from unittest.mock import patch + +from tests.runtime.run_runtime_validation import ( + FixtureConfig, + ModelRequirement, + ValidationFailure, + check_models, + ensure_runtime_path, + execute_generated_python, + load_fixture_names, + parse_png_info, + validate_generated_python, +) + + +def make_png_bytes( + width: int, + height: int, + text_chunks: list[tuple[bytes, bytes]] | None = None, +) -> bytes: + def chunk(chunk_type: bytes, data: bytes) -> bytes: + crc = zlib.crc32(chunk_type + data) & 0xFFFFFFFF + return ( + struct.pack(">I", len(data)) + + chunk_type + + data + + struct.pack(">I", crc) + ) + + ihdr = chunk(b"IHDR", struct.pack(">IIBBBBB", width, height, 8, 2, 0, 0, 0)) + text_chunks = text_chunks or [] + idat = chunk( + b"IDAT", + zlib.compress(b"\x00" + (b"\x00\x00\x00" * width)), + ) + return b"".join( + [ + b"\x89PNG\r\n\x1a\n", + ihdr, + *[chunk(chunk_type, payload) for chunk_type, payload in text_chunks], + idat, + chunk(b"IEND", b""), + ] + ) + + +class RuntimeValidationHarnessTest(unittest.TestCase): + def test_load_fixture_names_all_returns_registered_names(self): + fixture_names = load_fixture_names("all") + + self.assertGreaterEqual(len(fixture_names), 4) + self.assertEqual(fixture_names[0], "upscale-model-loader") + self.assertIn("unsafe-kwargs", fixture_names) + + def test_load_fixture_names_rejects_unknown_fixture(self): + with self.assertRaises(ValidationFailure) as context: + load_fixture_names("missing-fixture") + + self.assertEqual(context.exception.classification, "fixture bug") + self.assertIn("Unknown fixture", context.exception.message) + + def test_ensure_runtime_path_fast_tier_returns_env_without_validation(self): + with patch.dict("os.environ", {"COMFYUI_PATH": "/does/not/exist"}, clear=False): + runtime_path = ensure_runtime_path("fast") + + self.assertEqual(runtime_path, "/does/not/exist") + + @patch("tests.runtime.run_runtime_validation.get_comfyui_path", return_value="") + def test_ensure_runtime_path_runtime_tier_requires_valid_checkout(self, _mock_path): + with self.assertRaises(ValidationFailure) as context: + ensure_runtime_path("runtime") + + self.assertEqual(context.exception.classification, "environment/setup failure") + self.assertIn("Could not find a valid ComfyUI checkout", context.exception.message) + + def test_check_models_returns_only_missing_requirements(self): + with tempfile.TemporaryDirectory() as tmpdir: + runtime_path = Path(tmpdir) + present_dir = runtime_path / "models" / "checkpoints" + present_dir.mkdir(parents=True) + (present_dir / "present.safetensors").write_text("ok", encoding="utf-8") + fixture = FixtureConfig( + name="model-checks", + path=Path("unused.json"), + model_requirements=( + ModelRequirement( + filename="present.safetensors", + relative_dir="models/checkpoints", + source_url="https://example.invalid/present", + ), + ModelRequirement( + filename="missing.safetensors", + relative_dir="models/checkpoints", + source_url="https://example.invalid/missing", + ), + ), + ) + + missing = check_models(fixture, str(runtime_path)) + + self.assertEqual([item.filename for item in missing], ["missing.safetensors"]) + + def test_validate_generated_python_accepts_valid_python(self): + validate_generated_python("value = 1\n", "valid-fixture") + + def test_validate_generated_python_reports_syntax_regression(self): + with self.assertRaises(ValidationFailure) as context: + validate_generated_python("def broken(:\n", "broken-fixture") + + self.assertEqual(context.exception.classification, "repo regression") + self.assertIn("broken-fixture", context.exception.message) + + def test_parse_png_info_reads_dimensions_and_text_chunks(self): + compressed_text = zlib.compress(b"workflow data") + png_bytes = make_png_bytes( + width=3, + height=2, + text_chunks=[ + (b"tEXt", b"prompt\x00hello"), + (b"zTXt", b"workflow\x00\x00" + compressed_text), + (b"iTXt", b"comment\x00\x00\x00\x00\x00unicode text"), + ], + ) + + with tempfile.TemporaryDirectory() as tmpdir: + image_path = Path(tmpdir) / "artifact.png" + image_path.write_bytes(png_bytes) + + width, height, metadata = parse_png_info(image_path) + + self.assertEqual((width, height), (3, 2)) + self.assertEqual(metadata["prompt"], "hello") + self.assertEqual(metadata["workflow"], "workflow data") + self.assertEqual(metadata["comment"], "unicode text") + + def test_parse_png_info_rejects_non_png_files(self): + with tempfile.TemporaryDirectory() as tmpdir: + image_path = Path(tmpdir) / "artifact.bin" + image_path.write_bytes(b"not-a-png") + + with self.assertRaises(ValidationFailure) as context: + parse_png_info(image_path) + + self.assertEqual(context.exception.classification, "environment/setup failure") + self.assertIn("Expected PNG output", context.exception.message) + + def test_parse_png_info_requires_dimensions(self): + with tempfile.TemporaryDirectory() as tmpdir: + image_path = Path(tmpdir) / "artifact.png" + image_path.write_bytes( + b"\x89PNG\r\n\x1a\n" + + struct.pack(">I", 0) + + b"IEND" + + struct.pack(">I", 0) + ) + + with self.assertRaises(ValidationFailure) as context: + parse_png_info(image_path) + + self.assertEqual(context.exception.classification, "environment/setup failure") + self.assertIn("Could not read PNG dimensions", context.exception.message) + + @patch("tests.runtime.run_runtime_validation.get_runtime_python", return_value="/usr/bin/python") + @patch("tests.runtime.run_runtime_validation.subprocess.run") + def test_execute_generated_python_classifies_missing_torch_as_environment_failure( + self, + mock_run, + _mock_runtime_python, + ): + mock_run.return_value.returncode = 1 + mock_run.return_value.stderr = "ModuleNotFoundError: No module named 'torch'" + mock_run.return_value.stdout = "" + fixture = FixtureConfig( + name="runtime-fixture", + path=Path("unused.json"), + runtime_capable=True, + ) + + with tempfile.TemporaryDirectory() as tmpdir: + (Path(tmpdir) / "output").mkdir() + with self.assertRaises(ValidationFailure) as context: + execute_generated_python("print('hello')\n", fixture, tmpdir) + + self.assertEqual(context.exception.classification, "environment/setup failure") + + @patch("tests.runtime.run_runtime_validation.get_runtime_python", return_value="/usr/bin/python") + @patch("tests.runtime.run_runtime_validation.subprocess.run") + def test_execute_generated_python_classifies_missing_files_as_environment_failure( + self, + mock_run, + _mock_runtime_python, + ): + mock_run.return_value.returncode = 1 + mock_run.return_value.stderr = "No such file or directory: missing.png" + mock_run.return_value.stdout = "" + fixture = FixtureConfig( + name="runtime-fixture", + path=Path("unused.json"), + runtime_capable=True, + ) + + with tempfile.TemporaryDirectory() as tmpdir: + (Path(tmpdir) / "output").mkdir() + with self.assertRaises(ValidationFailure) as context: + execute_generated_python("print('hello')\n", fixture, tmpdir) + + self.assertEqual(context.exception.classification, "environment/setup failure") + + @patch("tests.runtime.run_runtime_validation.get_runtime_python", return_value="/usr/bin/python") + @patch("tests.runtime.run_runtime_validation.subprocess.run") + def test_execute_generated_python_classifies_other_failures_as_repo_regression( + self, + mock_run, + _mock_runtime_python, + ): + mock_run.return_value.returncode = 1 + mock_run.return_value.stderr = "ValueError: broken workflow" + mock_run.return_value.stdout = "" + fixture = FixtureConfig( + name="runtime-fixture", + path=Path("unused.json"), + runtime_capable=True, + ) + + with tempfile.TemporaryDirectory() as tmpdir: + (Path(tmpdir) / "output").mkdir() + with self.assertRaises(ValidationFailure) as context: + execute_generated_python("print('hello')\n", fixture, tmpdir) + + self.assertEqual(context.exception.classification, "repo regression") + + @patch("tests.runtime.run_runtime_validation.get_runtime_python", return_value="/usr/bin/python") + @patch("tests.runtime.run_runtime_validation.subprocess.run") + def test_execute_generated_python_requires_fresh_matching_artifact( + self, + mock_run, + _mock_runtime_python, + ): + mock_run.return_value.returncode = 0 + mock_run.return_value.stderr = "" + mock_run.return_value.stdout = "" + fixture = FixtureConfig( + name="runtime-fixture", + path=Path("unused.json"), + runtime_capable=True, + filename_prefix="expected_prefix", + ) + + with tempfile.TemporaryDirectory() as tmpdir: + (Path(tmpdir) / "output").mkdir() + with self.assertRaises(ValidationFailure) as context: + execute_generated_python("print('hello')\n", fixture, tmpdir) + + self.assertEqual(context.exception.classification, "repo regression") + self.assertIn("did not produce a new output file", context.exception.message) + + @patch("tests.runtime.run_runtime_validation.get_runtime_python", return_value="/usr/bin/python") + @patch("tests.runtime.run_runtime_validation.validate_output_artifact") + @patch("tests.runtime.run_runtime_validation.subprocess.run") + def test_execute_generated_python_validates_newest_matching_artifact( + self, + mock_run, + mock_validate_output, + _mock_runtime_python, + ): + mock_run.return_value.returncode = 0 + mock_run.return_value.stderr = "" + mock_run.return_value.stdout = "" + fixture = FixtureConfig( + name="runtime-fixture", + path=Path("unused.json"), + runtime_capable=True, + filename_prefix="expected_prefix", + ) + + with tempfile.TemporaryDirectory() as tmpdir: + output_dir = Path(tmpdir) / "output" + output_dir.mkdir() + older = output_dir / "expected_prefix_00001_.png" + older.write_bytes(make_png_bytes(1, 1)) + with patch( + "tests.runtime.run_runtime_validation.subprocess.run", + side_effect=self._write_runtime_artifact(mock_run.return_value, output_dir), + ): + execute_generated_python("print('hello')\n", fixture, tmpdir) + + validated_path = mock_validate_output.call_args[0][1] + self.assertEqual(validated_path.name, "expected_prefix_00002_.png") + + @staticmethod + def _write_runtime_artifact(result, output_dir: Path): + def side_effect(*_args, **_kwargs): + (output_dir / "expected_prefix_00002_.png").write_bytes(make_png_bytes(2, 2)) + return result + + return side_effect + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_upscale_model_loader_export.py b/tests/test_upscale_model_loader_export.py index d711c4b..5101c20 100644 --- a/tests/test_upscale_model_loader_export.py +++ b/tests/test_upscale_model_loader_export.py @@ -115,6 +115,51 @@ class HiddenPromptSeedNode: class UpscaleModelLoaderExportTest(unittest.TestCase): + def test_top_level_module_preserves_exporter_entrypoints(self): + from comfyui_to_python import main + + self.assertTrue(callable(ComfyUItoPython)) + self.assertTrue(callable(run)) + self.assertTrue(callable(main)) + + def test_export_places_comfyui_bootstrap_before_torch_import(self): + workflow = { + "1": { + "class_type": "LoadImage", + "inputs": { + "image": "example.png", + }, + } + } + + output = StringIO() + ComfyUItoPython( + workflow=json.dumps(workflow), + output_file=output, + node_class_mappings={ + "LoadImage": LoadImage, + }, + ) + + generated = output.getvalue() + + self.assertIn("def bootstrap_comfyui_runtime()", generated) + self.assertIn("import comfy.options", generated) + self.assertIn("comfy.options.enable_args_parsing()", generated) + self.assertIn("import cuda_malloc", generated) + self.assertLess( + generated.index("def bootstrap_comfyui_runtime()"), + generated.index("import torch"), + ) + self.assertLess( + generated.index("bootstrap_comfyui_runtime()"), + generated.index("import torch"), + ) + self.assertLess( + generated.index("import cuda_malloc"), + generated.index("import torch"), + ) + def test_upscale_workflow_uses_direct_upscale_model_loader_init(self): workflow = { "1": { @@ -196,10 +241,11 @@ class UpscaleModelLoaderExportTest(unittest.TestCase): generated = output.getvalue() - self.assertIn('extra_pnginfo = {', generated) - self.assertIn('"workflow": json.loads(', generated) + self.assertIn("def build_extra_pnginfo()", generated) + self.assertIn('"workflow": {', generated) self.assertIn('"version": 0.4', generated) self.assertIn('"nodes": []', generated) + self.assertNotIn('"workflow": json.loads(', generated) self.assertNotIn('"source": "workflow_api"', generated) def test_export_without_frontend_workflow_leaves_png_workflow_metadata_absent(self): @@ -223,7 +269,8 @@ class UpscaleModelLoaderExportTest(unittest.TestCase): generated = output.getvalue() - self.assertIn("extra_pnginfo = None", generated) + self.assertIn("def build_extra_pnginfo()", generated) + self.assertIn("return None", generated) self.assertNotIn('"workflow": json.loads(', generated) def test_export_preserves_unique_variable_names_for_subgraph_node_ids(self): @@ -296,9 +343,66 @@ class UpscaleModelLoaderExportTest(unittest.TestCase): generated = output_file.read_text(encoding="utf-8") - self.assertIn("extra_pnginfo = None", generated) + self.assertIn("def build_extra_pnginfo()", generated) + self.assertIn("return None", generated) self.assertNotIn('"workflow": json.loads(', generated) + def test_export_structures_generated_script_into_readable_sections(self): + workflow = { + "1": { + "class_type": "LoadImage", + "inputs": { + "image": "example.png", + }, + }, + "2": { + "class_type": "UpscaleModelLoader", + "inputs": { + "model_name": "RealESRGAN_x4plus.safetensors", + }, + }, + "3": { + "class_type": "ImageUpscaleWithModel", + "inputs": { + "upscale_model": ["2", 0], + "image": ["1", 0], + }, + }, + } + + output = StringIO() + ComfyUItoPython( + workflow=json.dumps(workflow), + output_file=output, + node_class_mappings={ + "LoadImage": LoadImage, + "UpscaleModelLoader": UpscaleModelLoader, + "ImageUpscaleWithModel": ImageUpscaleWithModel, + }, + ) + + generated = output.getvalue() + + self.assertIn("# Imports", generated) + self.assertIn("# Runtime support", generated) + self.assertIn("# Workflow data", generated) + self.assertIn("# Workflow execution", generated) + self.assertIn("# Entrypoint", generated) + self.assertIn("def build_workflow()", generated) + self.assertIn("def build_extra_pnginfo()", generated) + self.assertIn("def main()", generated) + self.assertNotIn("def initialize_workflow()", generated) + self.assertNotIn("def run_once(", generated) + self.assertIn("with torch.inference_mode():", generated) + self.assertIn("for q in range(1):", generated) + self.assertIn("workflow = build_workflow()", generated) + self.assertIn("extra_pnginfo = build_extra_pnginfo()", generated) + self.assertNotIn('workflow = json.loads("', generated) + self.assertLess( + generated.index("def build_workflow()"), + generated.index("def main()"), + ) + def test_hidden_metadata_kwargs_follow_function_signature(self): workflow = { "1": {