"""Tests for wp_nodes/context_loop.py. Covers the recovery-friendly `_parse_config` helper + the `WPContextLoop` execute path (single-iteration, multi-iteration, bypass, override on/off, iteration var stamping). """ import json import pytest from wp_nodes.context_loop import WPContextLoop, _parse_config _SWEEP_OFF = {"enabled": False, "limit": 64, "hold_others": True, "axes": []} _DEFAULTS = { "strategy": "hash_index", "override_seed": False, "iteration_var_name": "iteration", "bypass": False, "iteration_internal": True, "total_internal": True, "seed_locks": {}, "bypass_frames": [], "sweep": _SWEEP_OFF, } @pytest.mark.parametrize( "raw,expected", [ ("", _DEFAULTS), ("{}", _DEFAULTS), ("not json", _DEFAULTS), ("[1, 2, 3]", _DEFAULTS), ( '{"strategy": "sequential", "override_seed": true, ' '"iteration_var_name": "idx", "bypass": true}', {"strategy": "sequential", "override_seed": True, "iteration_var_name": "idx", "bypass": True, "iteration_internal": True, "total_internal": True, "seed_locks": {}, "bypass_frames": [], "sweep": _SWEEP_OFF}, ), ( '{"strategy": "wat", "override_seed": true}', {"strategy": "hash_index", "override_seed": True, "iteration_var_name": "iteration", "bypass": False, "iteration_internal": True, "total_internal": True, "seed_locks": {}, "bypass_frames": [], "sweep": _SWEEP_OFF}, ), ('{"iteration_var_name": " "}', _DEFAULTS), ( '{"iteration_internal": true, "total_internal": true}', {**_DEFAULTS, "iteration_internal": True, "total_internal": True}, ), ], ) def test_parse_config(raw, expected): assert _parse_config(raw) == expected def _execute(seed, count, config_dict): """Call WPContextLoop.execute and unwrap the NodeOutput. Stub `NodeOutput.values` is a tuple of positional args; real ComfyUI runtime uses the same surface. The first element is our payload list. """ config_json = json.dumps(config_dict) result = WPContextLoop.execute(seed=seed, count=count, wp_context_loop_config=config_json) if hasattr(result, "values"): return result.values[0] return result def test_execute_count_one_emits_single_payload(): payloads = _execute(42, 1, {"strategy": "hash_index"}) assert isinstance(payloads, list) assert len(payloads) == 1 assert payloads[0].internals["__wp_loop_index__"] == 0 assert payloads[0].internals["__wp_seed_override__"] is None def test_execute_count_three_emits_three_payloads(): payloads = _execute(42, 3, {"strategy": "sequential"}) assert len(payloads) == 3 for i, p in enumerate(payloads): assert p.internals["__wp_loop_index__"] == i def test_execute_bypass_collapses_to_one_iteration(): payloads = _execute(42, 5, {"bypass": True}) assert len(payloads) == 1 def test_execute_override_off_leaves_seed_override_none(): payloads = _execute(42, 3, {"override_seed": False, "strategy": "sequential"}) assert all(p.internals["__wp_seed_override__"] is None for p in payloads) assert all("__wp_loop_seeds__" not in p.internals for p in payloads) def test_execute_override_on_stamps_derived_series(): payloads = _execute(42, 3, {"override_seed": True, "strategy": "sequential"}) assert payloads[0].internals["__wp_seed_override__"] == 42 assert payloads[1].internals["__wp_seed_override__"] == 43 assert payloads[2].internals["__wp_seed_override__"] == 44 assert payloads[0].internals["__wp_loop_seeds__"] == [42, 43, 44] def test_execute_iteration_var_default(): # 1-based — user-facing counter reads naturally in templates. # Internal `__wp_loop_index__` stays 0-based (engine seed math). payloads = _execute(42, 2, {}) assert payloads[0].context == {"iteration": "1", "iteration_total": "2"} assert payloads[1].context == {"iteration": "2", "iteration_total": "2"} def test_execute_iteration_var_custom_name(): payloads = _execute(42, 2, {"iteration_var_name": "idx"}) assert payloads[0].context == {"idx": "1", "idx_total": "2"} assert payloads[1].context == {"idx": "2", "idx_total": "2"} def test_execute_count_clamps_to_one_min(): payloads = _execute(42, 0, {}) assert len(payloads) == 1 def test_execute_no_internal_flags_when_off(): # Explicitly opt both flags out — asserts no internal_flags entry. payloads = _execute(42, 2, { "iteration_var_name": "idx", "iteration_internal": False, "total_internal": False, }) for p in payloads: assert "__wp_internal_flags__" not in p.internals def test_execute_default_config_internal_flags_both_set(): # Empty config uses defaults: both iteration + total are internal. payloads = _execute(42, 2, {}) for p in payloads: flags = p.internals["__wp_internal_flags__"] assert flags == {"iteration": True, "iteration_total": True} def test_execute_iteration_internal_stamps_flag(): # Opt total out so only iteration appears in flags. payloads = _execute(42, 2, { "iteration_var_name": "idx", "iteration_internal": True, "total_internal": False, }) for p in payloads: flags = p.internals["__wp_internal_flags__"] assert flags == {"idx": True} def test_execute_total_internal_stamps_flag(): # Opt iteration out so only total appears in flags. payloads = _execute(42, 2, { "iteration_var_name": "idx", "iteration_internal": False, "total_internal": True, }) for p in payloads: flags = p.internals["__wp_internal_flags__"] assert flags == {"idx_total": True} def test_execute_both_internal_flags_stamped(): payloads = _execute(42, 2, { "iteration_var_name": "iteration", "iteration_internal": True, "total_internal": True, }) for p in payloads: flags = p.internals["__wp_internal_flags__"] assert flags == {"iteration": True, "iteration_total": True} # --------------------------------------------- loop_config side output --- def _execute_full(seed, count, config_dict): """Same as ``_execute`` but returns the full NodeOutput.values tuple.""" config_json = json.dumps(config_dict) result = WPContextLoop.execute(seed=seed, count=count, wp_context_loop_config=config_json) return result.values def test_execute_emits_loop_config_payload_on_second_output(): """The second output is the resolved config dict — count / strategy / base_seed / override_seed — for WP_SeedList + friends to mirror the loop's series.""" values = _execute_full(42, 3, {"strategy": "sequential", "override_seed": True}) assert len(values) == 2 loop_config = values[1] assert loop_config == { "count": 3, "active_count": 3, "kept_indices": [0, 1, 2], "strategy": "sequential", "base_seed": 42, "override_seed": True, } def test_loop_config_count_reflects_bypass_collapse(): """Bypass collapses to effective_count=1, and the side output must reflect what the loop is ACTUALLY running, not the raw widget.""" values = _execute_full(42, 5, {"bypass": True}) loop_config = values[1] assert loop_config["count"] == 1 def test_loop_config_strategy_falls_back_to_default_on_unknown(): """_parse_config sanitises unknown strategies to the default — the side output reflects the resolved strategy.""" values = _execute_full(42, 3, {"strategy": "wat"}) loop_config = values[1] assert loop_config["strategy"] == "hash_index" def test_wired_values_set_the_frames_and_write_iteration_value(): out = WPContextLoop.execute( seed=[7], count=[3], wp_context_loop_config=["{}"], values=["wolf\n\n fox ", "owl"], ) payloads, loop_cfg = out.values assert [p.context["iteration_value"] for p in payloads] == ["wolf", "fox", "owl"] assert [p.context["iteration_total"] for p in payloads] == ["3", "3", "3"] assert loop_cfg["count"] == 3 and loop_cfg["values_count"] == 3 assert out.ui["loop_values"] == [{"count": 3, "frames": 3}] # not internal: the value is meant to print assert "iteration_value" not in payloads[0].internals.get("__wp_internal_flags__", {}) def test_wired_values_times_sweep_combinations(): sweep = {"enabled": True, "axes": [{"uid": "u1", "option_ids": ["a", "b"]}]} out = WPContextLoop.execute( seed=0, count=1, wp_context_loop_config=json.dumps({"sweep": sweep}), values="x\ny", ) payloads = out.values[0] assert [p.context["iteration_value"] for p in payloads] == ["x", "x", "y", "y"] assert [p.internals["__wp_pin_overrides__"]["u1"] for p in payloads] == ["a", "b", "a", "b"] def test_empty_values_wire_keeps_the_count_and_bypass_reads_the_first(): out = WPContextLoop.execute(seed=0, count=2, wp_context_loop_config="{}", values=[" \n"]) assert len(out.values[0]) == 2 assert "iteration_value" not in out.values[0][0].context assert "loop_values" not in out.ui out = WPContextLoop.execute( seed=0, count=2, wp_context_loop_config=json.dumps({"bypass": True}), values="a\nb", ) assert [p.context["iteration_value"] for p in out.values[0]] == ["a"] def test_wired_values_cap_at_the_frame_limit(): out = WPContextLoop.execute( seed=0, count=1, wp_context_loop_config="{}", values="\n".join(map(str, range(1200))), ) assert len(out.values[0]) == 999 def test_frames_carry_debug_trace_rows(): """WP Debug names the Loop as the writer of its variables and shows the frame each payload is, with its seed when the Loop drives seeds.""" payloads = _execute(7, 4, { "override_seed": True, "bypass_frames": [1], "iteration_internal": True, "total_internal": False, }) assert [p.internals["__wp_loop_index__"] for p in payloads] == [0, 2, 3] rows = payloads[1].debug["__wp_trace__"] assert [(r["binding"], r["value"], r["internal"]) for r in rows] == [ ("iteration", "3", True), ("iteration_total", "4", False), ] assert all(r["node"] == "WP_ContextLoop" for r in rows) detail = rows[0]["detail"] assert detail["frame"] == 3 and detail["of"] == 4 assert detail["seed"] == payloads[1].internals["__wp_seed_override__"] def test_trace_rows_note_the_value_line(): out = WPContextLoop.execute( seed=0, count=1, wp_context_loop_config="{}", values=["a\nb"], ) payloads = out.values[0] rows = payloads[1].debug["__wp_trace__"] value_row = next(r for r in rows if r["binding"] == "iteration_value") assert value_row["value"] == "b" assert value_row["detail"]["value_line"] == 2 assert value_row["detail"]["values"] == 2 assert "seed" not in value_row["detail"]