"""Smoke test for the descriptor-driven ``local_validate.image_max_batch`` client-side fail-fast. Exercises the proxy_node's ``_collect_local_validate`` / ``_enforce_local_validate`` pair against a fake descriptor that mirrors the Flux2 wire shape. Asserts that: 1. ``_collect_local_validate`` extracts ``{"images": {"image_max_batch": 2}}`` from the descriptor's IMAGE-input options dict. 2. Encoding a 3-frame tensor for that input raises ``RnpProtocolError`` with ``code=INPUT_INVALID`` and ``user_facing=True`` *before* any serialization / HTTP attempt. 3. A 1-frame tensor passes through cleanly. 4. A descriptor without ``local_validate`` is a no-op. 5. ``cap=1`` surfaces the single-frame copy. 6. A ``local_validate`` block missing the ``image_max_batch`` key is a no-op. Run with any python that has ``torch`` available: python notes/run_image_max_batch.py The test loads the helper functions out of ``proxy_node.py`` via AST extraction so it doesn't pay the cost of importing ComfyUI's ``comfy_api`` / ``comfy_api_nodes`` packages (which the helper functions don't touch). Standalone — no test runner / fixtures needed. """ from __future__ import annotations import ast import asyncio import importlib.util import os import sys import types import torch # --------------------------------------------------------------------------- # Locate the worktree under test (this file's parent's parent). # --------------------------------------------------------------------------- _HERE = os.path.dirname(os.path.abspath(__file__)) _CLIENT_DIR = os.path.dirname(_HERE) def _load_protocol() -> types.ModuleType: """Load ``protocol.py`` standalone (no relative imports).""" path = os.path.join(_CLIENT_DIR, "protocol.py") spec = importlib.util.spec_from_file_location("rnp_smoke_protocol", path) mod = importlib.util.module_from_spec(spec) sys.modules[spec.name] = mod spec.loader.exec_module(mod) return mod def _build_serialization_stub(protocol: types.ModuleType) -> types.ModuleType: """Minimal ``serialization`` shim covering exactly the surface ``_encode_one`` / ``_enforce_local_validate`` / ``_encode_inputs`` use: * ``_is_torch_tensor(value)`` — duck-type check for torch tensors. * ``is_audio_input(value)`` — always False (no audio in this test). * ``encode_image_tensor(tensor, *, accepts_batch)`` — fake envelope. * ``encode_mask_tensor(tensor)`` — fake envelope. * ``is_envelope(value)`` — False so ``maybe_externalize`` is skipped. """ mod = types.ModuleType("rnp_smoke_serialization") def _is_torch_tensor(value): # noqa: ANN001 try: import torch as _torch except ImportError: return False return isinstance(value, _torch.Tensor) def is_audio_input(value): # noqa: ANN001 return isinstance(value, dict) and {"waveform", "sample_rate"} <= set(value.keys()) def encode_image_tensor(tensor, *, accepts_batch=False): # noqa: ANN001, ARG001 return {"type": "image", "encoding": "png_base64", "frames": int(tensor.shape[0])} def encode_mask_tensor(tensor): # noqa: ANN001, ARG001 return {"type": "mask", "encoding": "png_base64"} def encode_audio_input(value): # noqa: ANN001, ARG001 return {"type": "audio", "encoding": "mp3_base64"} def is_envelope(value): # noqa: ANN001 return False # smoke test never externalizes mod._is_torch_tensor = _is_torch_tensor mod.is_audio_input = is_audio_input mod.encode_image_tensor = encode_image_tensor mod.encode_mask_tensor = encode_mask_tensor mod.encode_audio_input = encode_audio_input mod.is_envelope = is_envelope async def maybe_externalize(*args, **kwargs): # pragma: no cover — unused return args[0] if args else None mod.maybe_externalize = maybe_externalize return mod def _build_proxy_node_stub( protocol: types.ModuleType, serialization: types.ModuleType, ) -> types.ModuleType: """Extract the four helper functions under test from ``proxy_node.py`` and exec them in a minimal namespace so we don't drag the rest of the module's heavy ComfyUI imports into the test process. """ mod = types.ModuleType("rnp_smoke_proxy_node") mod.__dict__.update({ "Any": object, "log": types.SimpleNamespace(warning=lambda *a, **kw: None), "serialization": serialization, "RnpProtocolError": protocol.RnpProtocolError, "ErrorCode": protocol.ErrorCode, }) src_path = os.path.join(_CLIENT_DIR, "proxy_node.py") with open(src_path, "r", encoding="utf-8") as fh: src = fh.read() tree = ast.parse(src) wanted = { "_collect_local_validate", "_rules_from_input_spec", "_enforce_local_validate", "_check_image_max_batch", "_encode_one", "_encode_inputs", } nodes = [ n for n in tree.body if isinstance(n, (ast.AsyncFunctionDef, ast.FunctionDef)) and n.name in wanted ] snippet = ast.Module(body=nodes, type_ignores=[]) code = compile(snippet, src_path, "exec") exec(code, mod.__dict__) missing = wanted - set(mod.__dict__) if missing: raise RuntimeError(f"proxy_node helpers missing from extract: {missing}") return mod # --------------------------------------------------------------------------- # Test fixtures # --------------------------------------------------------------------------- def _flux2_like_descriptor(cap: int) -> dict: """Mirrors the Flux2 ``images`` input descriptor surface.""" return { "name": "FakeFlux2_RNP", "input": { "required": { "prompt": ["STRING", {"multiline": True, "default": ""}], }, "optional": { "images": ["IMAGE", { "tooltip": "Up to N images to be used as references.", "local_validate": {"image_max_batch": cap}, }], }, "hidden": {}, }, } def _no_validate_descriptor() -> dict: return { "name": "FakeNoValidate_RNP", "input": { "required": {}, "optional": { "images": ["IMAGE", {"tooltip": "no cap"}], }, "hidden": {}, }, } def _autogrow_descriptor(cap: int, prefix: str = "reference") -> dict: """Mirrors GrokVideoReferenceNode_RNP's top-level Autogrow input after server descriptor unwraps the nested DynamicCombo branch (kept simpler here for direct AUTOGROW coverage).""" return { "name": "FakeAutogrow_RNP", "input": { "required": { "reference_images": ["AUTOGROW", { "template": ["IMAGE", { "local_validate": {"image_max_batch": cap}, }], "prefix": prefix, "min": 1, "max": cap, }], }, "optional": {}, "hidden": {}, }, } def _dyncombo_with_autogrow_descriptor() -> dict: """Mirrors GrokImageEditNodeV2_RNP's ``model`` DynamicCombo with two branches each carrying an Autogrow ``images`` template — one capped at 3, the other at 1.""" def branch(key: str, cap: int) -> dict: return { "key": key, "inputs": [ ["images", "AUTOGROW", { "template": ["IMAGE", { "local_validate": {"image_max_batch": cap}, }], "prefix": "image", "min": 1, "max": cap, }], ["resolution", "COMBO", {"options": ["1K", "2K"]}], ], } return { "name": "FakeDynComboAutogrow_RNP", "input": { "required": { "model": ["DYNAMIC_COMBO", { "options": [ branch("grok-imagine-image-quality", 3), branch("grok-imagine-image-pro", 1), ], }], }, "optional": {}, "hidden": {}, }, } # --------------------------------------------------------------------------- # Tests # --------------------------------------------------------------------------- async def main() -> int: protocol = _load_protocol() serialization = _build_serialization_stub(protocol) proxy_node = _build_proxy_node_stub(protocol, serialization) # ---- 1. _collect_local_validate extracts the rule desc = _flux2_like_descriptor(cap=2) rules = proxy_node._collect_local_validate(desc) assert rules == {"images": {"image_max_batch": 2}}, rules print("ok: _collect_local_validate extracted", rules) # ---- 2. over-cap tensor raises INPUT_INVALID before encoding over_cap = torch.zeros((3, 64, 64, 3), dtype=torch.float32) raised = False try: await proxy_node._encode_inputs( {"images": over_cap}, input_serialization={"images": ["png_base64", "png_base64_batch"]}, local_validate=rules, node_id="FakeFlux2_RNP", ) except protocol.RnpProtocolError as e: raised = True assert e.code == protocol.ErrorCode.INPUT_INVALID, e.code assert e.user_facing is True, e.user_facing msg = str(e) assert "FakeFlux2_RNP" in msg, msg assert "3" in msg and "2" in msg, msg print("ok: over-cap raised INPUT_INVALID:", msg) assert raised, "over-cap tensor did not raise" # ---- 3. at-cap tensor passes through (no exception) at_cap = torch.zeros((2, 64, 64, 3), dtype=torch.float32) out = await proxy_node._encode_inputs( {"images": at_cap}, input_serialization={"images": ["png_base64", "png_base64_batch"]}, local_validate=rules, node_id="FakeFlux2_RNP", ) assert "images" in out print("ok: at-cap (B=2) encoded without error") # ---- 4. cap=1 / multi-frame raises with the single-frame copy desc1 = _flux2_like_descriptor(cap=1) rules1 = proxy_node._collect_local_validate(desc1) raised = False try: await proxy_node._encode_inputs( {"images": over_cap}, local_validate=rules1, node_id="FakeI2V_RNP", ) except protocol.RnpProtocolError as e: raised = True assert e.code == protocol.ErrorCode.INPUT_INVALID assert "single image" in str(e), str(e) print("ok: cap=1 single-frame copy fired:", str(e)) assert raised # ---- 5. descriptor without local_validate is a no-op desc_n = _no_validate_descriptor() rules_n = proxy_node._collect_local_validate(desc_n) assert rules_n == {}, rules_n out = await proxy_node._encode_inputs( {"images": over_cap}, input_serialization={"images": ["png_base64", "png_base64_batch"]}, local_validate=rules_n, node_id="FakeNoValidate_RNP", ) assert "images" in out print("ok: missing local_validate is a no-op") # ---- 6. local_validate block without image_max_batch sub-key is a no-op out = await proxy_node._encode_inputs( {"images": over_cap}, input_serialization={"images": ["png_base64", "png_base64_batch"]}, local_validate={"images": {"some_other_rule": 7}}, node_id="FakeOther_RNP", ) assert "images" in out print("ok: missing image_max_batch sub-key is a no-op") # ---- 7. AUTOGROW: each slot value enforced against the template cap desc_ag = _autogrow_descriptor(cap=2) rules_ag = proxy_node._collect_local_validate(desc_ag) assert "reference_images" in rules_ag, rules_ag assert rules_ag["reference_images"].get("__template__") == {"image_max_batch": 2}, rules_ag print("ok: AUTOGROW collected __template__ rules:", rules_ag) one_frame = torch.zeros((1, 64, 64, 3), dtype=torch.float32) raised = False try: # Slot 1 is over-cap (B=3 > cap=2) → should fire on that slot await proxy_node._encode_inputs( {"reference_images": {"reference0": one_frame, "reference1": over_cap}}, local_validate=rules_ag, node_id="FakeAutogrow_RNP", ) except protocol.RnpProtocolError as e: raised = True assert e.code == protocol.ErrorCode.INPUT_INVALID msg = str(e) # Slot path label should mention reference_images. assert "reference_images" in msg, msg assert "reference1" in msg, msg print("ok: AUTOGROW slot over-cap fired:", msg) assert raised, "AUTOGROW slot over-cap did not raise" # All-under-cap autogrow passes (no exception) out = await proxy_node._encode_inputs( {"reference_images": {"reference0": one_frame, "reference1": one_frame, "reference2": None}}, local_validate=rules_ag, node_id="FakeAutogrow_RNP", ) assert "reference_images" in out print("ok: AUTOGROW all-under-cap (None slots tolerated) passed") # ---- 8. DYNAMIC_COMBO with nested AUTOGROW: branch-aware enforcement desc_dc = _dyncombo_with_autogrow_descriptor() rules_dc = proxy_node._collect_local_validate(desc_dc) assert "model" in rules_dc, rules_dc branches = rules_dc["model"].get("__branches__") assert isinstance(branches, dict), rules_dc assert "grok-imagine-image-quality" in branches assert "grok-imagine-image-pro" in branches assert branches["grok-imagine-image-quality"]["images"]["__template__"] == {"image_max_batch": 3} assert branches["grok-imagine-image-pro"]["images"]["__template__"] == {"image_max_batch": 1} print("ok: DYNAMIC_COMBO collected branch-keyed templates") # Branch A (cap=3): slot with B=4 should raise over_4 = torch.zeros((4, 64, 64, 3), dtype=torch.float32) raised = False try: await proxy_node._encode_inputs( {"model": { "model": "grok-imagine-image-quality", "resolution": "1K", "images": {"image0": over_4}, }}, local_validate=rules_dc, node_id="FakeDynComboAutogrow_RNP", ) except protocol.RnpProtocolError as e: raised = True assert e.code == protocol.ErrorCode.INPUT_INVALID msg = str(e) assert "image0" in msg and "3" in msg and "4" in msg, msg print("ok: DYNAMIC_COMBO branch-A (cap=3) fired on B=4 slot:", msg) assert raised # Branch A (cap=3): all slots within cap → passes out = await proxy_node._encode_inputs( {"model": { "model": "grok-imagine-image-quality", "resolution": "2K", "images": {"image0": one_frame, "image1": one_frame, "image2": one_frame}, }}, local_validate=rules_dc, node_id="FakeDynComboAutogrow_RNP", ) assert "model" in out print("ok: DYNAMIC_COMBO branch-A all-under-cap passed") # Branch B (cap=1): a B=3 slot fires the single-frame copy raised = False try: await proxy_node._encode_inputs( {"model": { "model": "grok-imagine-image-pro", "resolution": "2K", "images": {"image0": over_cap}, }}, local_validate=rules_dc, node_id="FakeDynComboAutogrow_RNP", ) except protocol.RnpProtocolError as e: raised = True assert e.code == protocol.ErrorCode.INPUT_INVALID assert "single image" in str(e), str(e) print("ok: DYNAMIC_COMBO branch-B (cap=1) fired single-frame copy:", str(e)) assert raised # Unknown branch key → no enforcement (defensive) out = await proxy_node._encode_inputs( {"model": { "model": "unknown-branch-key", "images": {"image0": over_cap}, }}, local_validate=rules_dc, node_id="FakeDynComboAutogrow_RNP", ) assert "model" in out print("ok: DYNAMIC_COMBO unknown branch key is a defensive no-op") print("PASS: image_max_batch local_validate smoke test (incl. nested traversal)") return 0 if __name__ == "__main__": sys.exit(asyncio.run(main()))