Strengthen test suite robustness

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
 

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
+      ]
+    }
+  }
+}

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
+      ]
+    }
+  }
+}

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()

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()

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,
+        },
+    ),
 }
 
 

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()

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": {