aigc-home/calliope-backend/tests/test_ui_convert.py

"""HEL-681: UI→API converter unit tests + ground-truth cross-checks.

The synthetic graphs pin the structural semantics (reroute / bypass / mute /
subgraph expansion / Set-Get variables / control_after_generate pairing);
the corpus section runs the real 146-file library and the repo's API exports.
"""

from __future__ import annotations

import copy
import json
from pathlib import Path

import pytest

from calliope.comfyui.lint import workflow_warnings
from calliope.comfyui.ui_convert import convert_ui_workflow

EXAMPLES = Path(__file__).resolve().parents[2] / "example_ComfyUI_workflows"


def _node(nid, ntype, *, mode=0, inputs=None, outputs=None, wv=None, wvn=None, title=None):
    node = {"id": nid, "type": ntype, "mode": mode, "inputs": inputs or [], "order": 0}
    if outputs is not None:
        node["outputs"] = outputs
    if wv is not None:
        node["widgets_values"] = wv
    if wvn is not None:
        node["widgets_values_named"] = wvn
    if title:
        node["title"] = title
    return node


def _in_slot(name, type_, link=None, widget=None):
    entry = {"name": name, "type": type_, "link": link}
    if widget:
        entry["widget"] = {"name": widget}
    return entry


def _out(name, type_, links=None):
    return {"name": name, "type": type_, "links": links or [], "slot_index": 0}


def _base_graph():
    """text2image: loader → prompts/latent → KSampler → VAE → SaveImage."""
    nodes = [
        _node(
            1,
            "CheckpointLoaderSimple",
            outputs=[_out("MODEL", "MODEL"), _out("CLIP", "CLIP"), _out("VAE", "VAE")],
            wv=["model.safetensors"],
            inputs=[_in_slot("ckpt_name", "COMBO", widget="ckpt_name")],
        ),
        _node(
            2,
            "CLIPTextEncode",
            inputs=[
                _in_slot("clip", "CLIP", link=1),
                _in_slot("text", "STRING", widget="text"),
            ],
            outputs=[_out("CONDITIONING", "CONDITIONING")],
            wv=["a cat"],
        ),
        _node(
            3,
            "CLIPTextEncode",
            inputs=[
                _in_slot("clip", "CLIP", link=2),
                _in_slot("text", "STRING", widget="text"),
            ],
            outputs=[_out("CONDITIONING", "CONDITIONING")],
            wv=["bad quality"],
        ),
        _node(
            4,
            "EmptyLatentImage",
            inputs=[
                _in_slot("width", "INT", widget="width"),
                _in_slot("height", "INT", widget="height"),
                _in_slot("batch_size", "INT", widget="batch_size"),
            ],
            outputs=[_out("LATENT", "LATENT")],
            wv=[512, 512, 1],
        ),
        _node(
            5,
            "KSampler",
            inputs=[
                _in_slot("model", "MODEL", link=3),
                _in_slot("positive", "CONDITIONING", link=4),
                _in_slot("negative", "CONDITIONING", link=5),
                _in_slot("latent_image", "LATENT", link=6),
                _in_slot("seed", "INT", widget="seed"),
                _in_slot("steps", "INT", widget="steps"),
                _in_slot("cfg", "FLOAT", widget="cfg"),
                _in_slot("sampler_name", "COMBO", widget="sampler_name"),
                _in_slot("scheduler", "COMBO", widget="scheduler"),
                _in_slot("denoise", "FLOAT", widget="denoise"),
            ],
            outputs=[_out("LATENT", "LATENT")],
            # seed 42 + control "fixed" rides together, then five widgets
            wv=[42, "fixed", 20, 7, "euler", "normal", 1.0],
        ),
        _node(
            6,
            "VAEDecode",
            inputs=[
                _in_slot("samples", "LATENT", link=7),
                _in_slot("vae", "VAE", link=8),
            ],
            outputs=[_out("IMAGE", "IMAGE")],
        ),
        _node(
            7,
            "SaveImage",
            inputs=[
                _in_slot("images", "IMAGE", link=9),
                _in_slot("filename_prefix", "STRING", widget="filename_prefix"),
            ],
            wv=["test"],
        ),
    ]
    links = [
        [1, 1, 1, 2, 0, "CLIP"],
        [2, 1, 1, 3, 0, "CLIP"],
        [3, 1, 0, 5, 0, "MODEL"],
        [4, 2, 0, 5, 1, "CONDITIONING"],
        [5, 3, 0, 5, 2, "CONDITIONING"],
        [6, 4, 0, 5, 3, "LATENT"],
        [7, 5, 0, 6, 0, "LATENT"],
        [8, 1, 2, 6, 1, "VAE"],
        [9, 6, 0, 7, 0, "IMAGE"],
    ]
    return {"nodes": nodes, "links": links, "last_node_id": 7, "last_link_id": 9}


def _convert(graph):
    return convert_ui_workflow(graph)


def test_clean_conversion_maps_everything():
    result = _convert(_base_graph())
    assert result.status == "ok"
    assert result.needs_confirmation is False
    wf = result.workflow
    assert set(wf) == {"1", "2", "3", "4", "5", "6", "7"}
    ks = wf["5"]
    assert ks["class_type"] == "KSampler"
    assert ks["inputs"]["seed"] == 42  # control_after_generate consumed silently
    assert ks["inputs"]["steps"] == 20
    assert ks["inputs"]["denoise"] == 1.0
    assert ks["inputs"]["model"] == ["1", 0]
    assert ks["inputs"]["latent_image"] == ["4", 0]
    # lint on the converted graph must be error-free
    assert not [w for w in workflow_warnings(wf) if w["level"] == "error"]


def test_upload_button_widgets_are_consumed_but_never_emitted():
    graph = _base_graph()
    load = _node(
        8,
        "LoadImage",
        inputs=[
            _in_slot("image", "COMBO", widget="image"),
            _in_slot("upload", "IMAGEUPLOAD", widget="upload"),
        ],
        outputs=[_out("IMAGE", "IMAGE")],
        wv=["sample.png", "image"],
    )
    save = _node(
        9,
        "PreviewImage",
        inputs=[
            _in_slot("images", "IMAGE", link=20),
            _in_slot("filename_prefix", "STRING", widget="filename_prefix"),
        ],
        wv=["preview"],
    )
    graph["nodes"] += [load, save]
    graph["links"].append([20, 8, 0, 9, 0, "IMAGE"])
    result = _convert(graph)
    assert result.workflow["8"]["inputs"] == {"image": "sample.png"}


def test_reroute_pass_through():
    graph = _base_graph()
    reroute = _node(8, "Reroute", inputs=[_in_slot("(in)", "*", link=9)])
    graph["nodes"].append(reroute)
    # SaveImage now fed by the reroute instead of VAEDecode directly
    graph["nodes"][6]["inputs"][0]["link"] = 10
    graph["links"].append([10, 8, 0, 7, 0, "IMAGE"])
    result = _convert(graph)
    assert result.workflow["7"]["inputs"]["images"] == ["6", 0]
    assert "8" not in result.workflow


def test_bypass_pass_through_and_mute_drop():
    graph = _base_graph()
    # bypass the KSampler: VAEDecode.samples must fall back to the latent
    graph["nodes"][4]["mode"] = 4
    result = _convert(graph)
    assert result.workflow["6"]["inputs"]["samples"] == ["4", 0]
    # mute the negative prompt: the link input is dropped with a warning
    graph = _base_graph()
    graph["nodes"][2]["mode"] = 2
    result = _convert(graph)
    assert "negative" not in result.workflow["5"]["inputs"]
    assert any(i.code == "muted_input_dropped" for i in result.issues)


def test_widget_entry_with_muted_upstream_falls_back_to_value():
    graph = _base_graph()
    # feed seed from a muted PrimitiveInt via a converted widget input
    prim = _node(
        9,
        "PrimitiveInt",
        mode=2,
        inputs=[_in_slot("value", "INT", widget="value"), _in_slot("control", "COMBO", widget="control")],
        wv=[7],
    )
    graph["nodes"].append(prim)
    graph["nodes"][4]["inputs"][4]["link"] = 11  # seed ← primitive
    graph["links"].append([11, 9, 0, 5, 4, "INT"])
    result = _convert(graph)
    assert result.workflow["5"]["inputs"]["seed"] == 42


def test_subgraph_expansion_and_boundary_wiring():
    graph = _base_graph()
    # pack the EmptyLatentImage into a subgraph exposing width/height widgets
    del graph["nodes"][3]
    graph["nodes"] = [n for n in graph["nodes"] if n["id"] != 4]
    sub_id = "aaaa-bbbb-cccc-dddd-eeee"
    packed = _node(
        20,
        sub_id,
        inputs=[
            _in_slot("width", "INT", widget="width"),
            _in_slot("height", "INT", widget="height"),
        ],
        outputs=[_out("LATENT", "LATENT")],
        wv=[640, 384],
    )
    graph["nodes"].append(packed)
    # KSampler latent_image ← packed output
    graph["links"] = [l for l in graph["links"] if l[0] != 6]
    graph["links"].append([6, 20, 0, 5, 3, "LATENT"])
    graph["definitions"] = {
        "subgraphs": [
            {
                "id": sub_id,
                "name": "Latent",
                "inputs": [
                    {"name": "width", "type": "INT"},
                    {"name": "height", "type": "INT"},
                ],
                "outputs": [{"name": "LATENT", "type": "LATENT"}],
                "inputNode": {"id": -10},
                "outputNode": {"id": -20},
                "nodes": [
                    _node(
                        5,
                        "EmptyLatentImage",
                        inputs=[
                            _in_slot("width", "INT", widget="width", link=101),
                            _in_slot("height", "INT", widget="height", link=102),
                            _in_slot("batch_size", "INT", widget="batch_size"),
                        ],
                        outputs=[_out("LATENT", "LATENT")],
                        wv=[512, 512, 1],
                    ),
                ],
                "links": [
                    {"id": 101, "origin_id": -10, "origin_slot": 0, "target_id": 5, "target_slot": 0, "type": "INT"},
                    {"id": 102, "origin_id": -10, "origin_slot": 1, "target_id": 5, "target_slot": 1, "type": "INT"},
                    {"id": 103, "origin_id": 5, "origin_slot": 0, "target_id": -20, "target_slot": 0, "type": "LATENT"},
                ],
            }
        ]
    }
    result = _convert(graph)
    wf = result.workflow
    assert "20:5" in wf  # qualified packed:inner id, ComfyUI's own scheme
    inner = wf["20:5"]
    assert inner["class_type"] == "EmptyLatentImage"
    # exposed widget values from the packed instance, not the inner defaults
    assert inner["inputs"]["width"] == 640
    assert inner["inputs"]["height"] == 384
    assert inner["inputs"]["batch_size"] == 1
    assert wf["5"]["inputs"]["latent_image"] == ["20:5", 0]
    # now wire the packed width slot from outside: boundary becomes a link
    res_node = _node(
        21,
        "ResolutionSelector",
        inputs=[_in_slot("aspect_ratio", "COMBO", widget="aspect_ratio")],
        outputs=[_out("width", "INT"), _out("height", "INT")],
        wv=["1:1"],
    )
    graph["nodes"].append(res_node)
    graph["nodes"][-2]["inputs"][0]["link"] = 12  # packed width ← selector
    graph["links"].append([12, 21, 0, 20, 0, "INT"])
    wf2 = _convert(graph).workflow
    assert wf2["20:5"]["inputs"]["width"] == ["21", 0]
    assert wf2["20:5"]["inputs"]["height"] == 384


def test_set_get_variables_resolve_to_captured_source():
    graph = _base_graph()
    setter = _node(8, "SetNode", inputs=[_in_slot("*", "*", link=8)], wv=["the_vae"])
    getter = _node(9, "GetNode", wv=["the_vae"], outputs=[_out("VAE", "VAE")])
    graph["nodes"] += [setter, getter]
    # VAEDecode.vae now fed via the variable
    graph["nodes"][5]["inputs"][1]["link"] = 13
    graph["links"] = [l for l in graph["links"] if l[0] != 8]
    graph["links"] += [
        [8, 1, 2, 8, 0, "VAE"],
        [13, 9, 0, 6, 1, "VAE"],
    ]
    result = _convert(graph)
    assert "8" not in result.workflow and "9" not in result.workflow
    assert result.workflow["6"]["inputs"]["vae"] == ["1", 2]


def test_unknown_widget_order_lands_in_needs_confirmation():
    graph = _base_graph()
    # an old-frontend node: positional widgets, no widget markers, no registry
    graph["nodes"].append(
        _node(10, "MysteryCustomNode", inputs=[], outputs=[_out("IMAGE", "IMAGE")], wv=[1, 2, 3])
    )
    save = _node(
        11,
        "PreviewImage",
        inputs=[_in_slot("images", "IMAGE", link=30)],
        wv=["preview"],
    )
    graph["nodes"].append(save)
    graph["links"].append([30, 10, 0, 11, 0, "IMAGE"])
    result = _convert(graph)
    assert result.needs_confirmation is True
    assert any(i.code == "widget_unmapped" for i in result.issues)
    # the node is still converted structurally — the report carries the doubt
    assert "10" in result.workflow


def test_not_ui_format_fails_fast():
    result = convert_ui_workflow({"3": {"class_type": "SaveImage", "inputs": {}}})
    assert result.status == "fail"
    assert result.workflow is None


# ---------------------------------------------------------------------------
# Real corpus: the 146 UI files + ground-truth API cross-checks (对拍).
# ---------------------------------------------------------------------------

pytestmark_convert = pytest.mark.skipif(
    not EXAMPLES.is_dir(), reason="example_ComfyUI_workflows not checked out"
)


@pytest.mark.skipif(not EXAMPLES.is_dir(), reason="example library missing")
def test_corpus_146_files_all_convert():
    ui_files = [
        p
        for p in sorted(EXAMPLES.rglob("*.json"))
        if isinstance(json.loads(p.read_text(encoding="utf-8")), dict)
        and isinstance(json.loads(p.read_text(encoding="utf-8")).get("nodes"), list)
    ]
    assert len(ui_files) == 146
    statuses = {"ok": 0, "warn": 0, "fail": 0}
    needs_confirmation = 0
    for path in ui_files:
        data = json.loads(path.read_text(encoding="utf-8"))
        result = convert_ui_workflow(data)
        statuses[result.status] += 1
        needs_confirmation += result.needs_confirmation
        assert result.workflow, f"{path.name} produced no workflow"
        # structural sanity: no dangling links in any conversion
        nodes = result.workflow
        for nid, node in nodes.items():
            for name, val in node.get("inputs", {}).items():
                if isinstance(val, list):
                    assert str(val[0]) in nodes, f"{path.name}: {nid}.{name} dangles"
    assert statuses["fail"] == 0
    assert statuses["ok"] + statuses["warn"] == 146
    # files whose widget order cannot be trusted are flagged, never hidden
    assert 0 < needs_confirmation < 10


@pytest.mark.skipif(not EXAMPLES.is_dir(), reason="example library missing")
def test_cross_check_against_ground_truth_exports():
    """The H3 t2v UI export converts to the same node classes the author's
    own API export carries; shared classes agree on scalar input names."""
    ui = json.loads(
        (EXAMPLES / "4-视频生成/Minimax-h3/minimax_h3_t2v_文生视频.json").read_text(encoding="utf-8")
    )
    api = json.loads(
        (EXAMPLES / "minimax_h3-Turbo_Text2Video_20260813_API.json").read_text(encoding="utf-8")
    )
    result = convert_ui_workflow(ui)
    assert result.status in ("ok", "warn")
    converted = result.workflow
    api_classes = {n["class_type"] for n in api.values()}
    # Every scalar name the conversion emits for a shared class must be an
    # input name the author's own export knows — wiring may differ between
    # versions (a muted branch falls back to its widget value), but the NAME
    # mapping must never invent or swap inputs.
    for cls in api_classes & {n["class_type"] for n in converted.values()}:
        api_names = {
            key
            for n in api.values()
            if n["class_type"] == cls
            for key in n.get("inputs", {})
        }
        conv_names = {
            key
            for n in converted.values()
            if n["class_type"] == cls
            for key, v in n.get("inputs", {}).items()
            if not isinstance(v, list)
        }
        assert conv_names <= api_names, f"class {cls} invents inputs: {conv_names - api_names}"


@pytest.mark.skipif(not EXAMPLES.is_dir(), reason="example library missing")
def test_packed_krea2_conversion_shape_matches_export_scheme():
    """The krea2 turbo file (packed subgraph) converts ok, expands inner nodes
    under `<packed>:<inner>` ids and maps the sampler widgets exactly — the
    same scheme the author's own Krea2 API export uses (ids like "171:153")."""
    ui = json.loads(
        (EXAMPLES / "3-图像生成/krea2/krea2_turbo_文生图.json").read_text(encoding="utf-8")
    )
    api = json.loads(
        (EXAMPLES / "Krea2_t2i_20260818_API.json").read_text(encoding="utf-8")
    )
    result = convert_ui_workflow(ui)
    assert result.status == "ok"
    assert any(":" in nid for nid in result.workflow)  # expanded subgraph nodes
    assert any("171:" in nid for nid in api)  # ground truth uses the same scheme
    ks = [
        n
        for n in result.workflow.values()
        if n["class_type"] == "KSampler" and not isinstance(n["inputs"].get("seed"), list)
    ]
    assert ks, "inner KSampler with widget-mapped seed expected"
    assert {"seed", "steps", "cfg", "sampler_name", "scheduler", "denoise"} <= set(
        ks[0]["inputs"]
    )