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

"""L2 连接核查（HEL-643 批次 1）：导入/复查时用目标 ComfyUI 的 /object_info
预检节点与模型可用性。全部基于 fake ComfyUI 的脚本化 /object_info，
不针对任何固定 JSON 或模型名写死判定。"""

from __future__ import annotations

import json
import socket
from pathlib import Path

import pytest

from calliope.comfyui.client import clear_object_info_cache
from calliope.comfyui.preflight import check_workflow_against
from calliope.config import settings
from tests.test_studio_generation import _url

FIXTURES = Path(__file__).parent / "fixtures" / "comfy_native"


def _load_05() -> dict:
    return json.loads((FIXTURES / "h3_ref2va_05_API.json").read_text(encoding="utf-8"))


def build_object_info(workflow: dict) -> dict:
    """A catalogue where every widget value the JSON carries is legal.

    String widgets become combos whose candidates include the workflow's own
    value (merged across nodes of the same class); numbers/booleans are typed
    inputs; wired inputs are link types — none of those can mismatch.
    """
    info: dict = {}
    for node in workflow.values():
        if not isinstance(node, dict):
            continue
        ct = str(node.get("class_type") or "")
        if not ct:
            continue
        spec = info.setdefault(ct, {"input": {"required": {}}, "output": ["*"]})
        required = spec["input"]["required"]
        for name, val in (node.get("inputs") or {}).items():
            if isinstance(val, str):
                # Merge: two nodes of one class with different values → both legal.
                if name in required and isinstance(required[name][0], list):
                    if val not in required[name][0]:
                        required[name][0] = [*required[name][0], val]
                    continue
                required[name] = [[val], {}]
            elif name in required:
                continue
            elif isinstance(val, list):
                required[name] = ["*", {}]
            elif isinstance(val, bool):
                required[name] = ["BOOLEAN", {}]
            elif isinstance(val, (int, float)):
                required[name] = ["NUMBER", {}]
    return info


@pytest.fixture(autouse=True)
def _fresh_cache():
    clear_object_info_cache()
    yield
    clear_object_info_cache()


@pytest.fixture
def endpoint_at(client, fake_comfy):  # noqa: F811
    """Register a connection pointing at the fake server; yields its id."""
    prev_eps = [dict(ep) for ep in (settings.comfyui_endpoints or [])]
    prev_active = settings.comfyui_active_id
    r = client.post(
        "/api/comfyui/endpoints",
        json={"name": "预检桩", "base_url": _url(fake_comfy), "auth_type": "none"},
    )
    assert r.status_code == 200, r.text
    endpoints = r.json()["endpoints"]
    ep = next(e for e in endpoints if e["base_url"] == _url(fake_comfy))
    settings.comfyui_endpoints = [dict(e) for e in endpoints if e["base_url"] == _url(fake_comfy)]
    settings.comfyui_active_id = ep["id"]
    settings.apply_active_comfyui()
    try:
        yield ep["id"]
    finally:
        settings.comfyui_endpoints = prev_eps
        settings.comfyui_active_id = prev_active
        settings.apply_active_comfyui()


# ---------------------------------------------------------------------------
# Unit rules
# ---------------------------------------------------------------------------


def test_unit_missing_node_reported_with_node_context():
    wf = {"900": {"class_type": "H3SLAAttention", "inputs": {"enabled": True},
                  "_meta": {"title": "H3 SLA Attention"}}}
    findings = check_workflow_against({"OtherNode": {}}, wf)
    assert len(findings) == 1
    f = findings[0]
    assert f["code"] == "missing_node" and f["level"] == "error"
    msg = f["message"]
    assert "H3SLAAttention" in msg and "900" in msg and "H3 SLA Attention" in msg
    assert "重新检查" in f["hint"]


def test_unit_combo_value_not_in_list():
    wf = {"127": {"class_type": "UNETLoader",
                  "inputs": {"unet_name": "missing_model.safetensors", "weight_dtype": "default"}}}
    info = {"UNETLoader": {"input": {"required": {
        "unet_name": [["installed_model.safetensors"], {}],
        "weight_dtype": [["default", "fp8"], {"upload": False}],
    }}, "output": ["MODEL"]}}
    findings = check_workflow_against(info, wf)
    assert len(findings) == 1
    f = findings[0]
    assert f["code"] == "value_not_in_list" and f["node_id"] == "127"
    assert "missing_model.safetensors" in f["message"] and "unet_name" in f["message"]
    assert "模型文件" in f["message"]
    assert "installed_model.safetensors" in f["hint"]


def test_unit_upload_backed_and_wired_inputs_skipped():
    wf = {
        "1": {"class_type": "LoadImage", "inputs": {"image": "placeholder.png"}},
        "2": {"class_type": "LoadVideo", "inputs": {"file": "clip.mp4"}},
        "3": {"class_type": "LoadAudio", "inputs": {"audio": "voice.wav"}},
        "4": {"class_type": "KSampler", "inputs": {"seed": 7, "model": ["9", 0]}},
    }
    info = {
        "LoadImage": {"input": {"required": {"image": [["other.png"], {"image_upload": True}]}},
                      "output": ["IMAGE"]},
        "LoadVideo": {"input": {"required": {"file": [[], {}]}}, "output": ["IMAGE"]},
        "LoadAudio": {"input": {"required": {"audio": [[], {}]}}, "output": ["AUDIO"]},
        "KSampler": {"input": {"required": {"seed": ["INT", {}], "model": ["MODEL", {}]}},
                     "output": ["IMAGE"]},
    }
    assert check_workflow_against(info, wf) == []


# ---------------------------------------------------------------------------
# Analyze with a target connection
# ---------------------------------------------------------------------------


def test_analyze_all_green_against_fake_object_info(client, endpoint_at, fake_comfy):
    fake_comfy.object_info = build_object_info(_load_05())
    r = client.post(
        "/api/workflows/analyze",
        json={"workflow_json": _load_05(), "endpoint_id": endpoint_at},
    )
    assert r.status_code == 200, r.text
    check = r.json()["connection_check"]
    assert check["reachable"] is True and check["ok"] is True
    assert check["findings"] == []


def test_analyze_reports_missing_custom_node(client, endpoint_at, fake_comfy):
    info = build_object_info(_load_05())
    info.pop("H3SLAAttention")
    fake_comfy.object_info = info
    r = client.post(
        "/api/workflows/analyze",
        json={"workflow_json": _load_05(), "endpoint_id": endpoint_at},
    )
    check = r.json()["connection_check"]
    assert check["ok"] is False
    missing = [f for f in check["findings"] if f["code"] == "missing_node"]
    assert missing and "H3SLAAttention" in missing[0]["message"]
    assert "900" in missing[0]["message"]


def test_analyze_reports_missing_model_value(client, endpoint_at, fake_comfy):
    info = build_object_info(_load_05())
    # The server has a DIFFERENT unet installed than the one the JSON names.
    info["UNETLoader"]["input"]["required"]["unet_name"] = [["another_unet.safetensors"], {}]
    fake_comfy.object_info = info
    r = client.post(
        "/api/workflows/analyze",
        json={"workflow_json": _load_05(), "endpoint_id": endpoint_at},
    )
    check = r.json()["connection_check"]
    assert check["ok"] is False
    bad = [f for f in check["findings"] if f["code"] == "value_not_in_list"]
    assert bad and bad[0]["node_id"] == "127"
    assert "minimax_h3_fused_refdelta" in bad[0]["message"]


def test_unreachable_connection_degrades_to_unchecked(client):
    # A port that nothing listens on: the endpoint resolves, the server doesn't.
    s = socket.socket()
    s.bind(("127.0.0.1", 0))
    port = s.getsockname()[1]
    s.close()
    prev_eps = [dict(ep) for ep in (settings.comfyui_endpoints or [])]
    prev_active = settings.comfyui_active_id
    r = client.post(
        "/api/comfyui/endpoints",
        json={"name": "断连桩", "base_url": f"http://127.0.0.1:{port}", "auth_type": "none"},
    )
    endpoints = r.json()["endpoints"]
    ep = next(e for e in endpoints if e["base_url"] == f"http://127.0.0.1:{port}")
    settings.comfyui_endpoints = [dict(e) for e in endpoints if e["id"] == ep["id"]]
    settings.comfyui_active_id = ep["id"]
    settings.apply_active_comfyui()
    try:
        r = client.post(
            "/api/workflows/analyze",
            json={"workflow_json": _load_05(), "endpoint_id": ep["id"]},
        )
        check = r.json()["connection_check"]
        assert check["reachable"] is False and check["ok"] is False
        assert "无法核查" in check["message"]
        assert check["findings"] == []
    finally:
        settings.comfyui_endpoints = prev_eps
        settings.comfyui_active_id = prev_active
        settings.apply_active_comfyui()


# ---------------------------------------------------------------------------
# Saved-workflow re-check + cache
# ---------------------------------------------------------------------------


def _create_workflow(client, workflow: dict, endpoint_id: str | None = None) -> int:
    r = client.post(
        "/api/workflows",
        json={
            "name": "预检工作流",
            "workflow_json": workflow,
            "usage": "ref2video",
            **({"endpoint_id": endpoint_id} if endpoint_id else {}),
        },
    )
    assert r.status_code == 200, r.text
    return r.json()["id"]


def test_saved_workflow_recheck_and_refresh(client, endpoint_at, fake_comfy):
    wf_id = _create_workflow(client, _load_05(), endpoint_at)
    info = build_object_info(_load_05())
    info.pop("UNETLoader")
    fake_comfy.object_info = info
    r = client.post(f"/api/workflows/{wf_id}/check", json={})
    assert r.status_code == 200, r.text
    body = r.json()
    assert body["ok"] is False and body["endpoint_id"] == endpoint_at
    assert any(f["code"] == "missing_node" for f in body["findings"])

    # User installs the node package, comes back and re-checks with refresh.
    fake_comfy.object_info = build_object_info(_load_05())
    r = client.post(f"/api/workflows/{wf_id}/check", json={"refresh": True})
    body = r.json()
    assert body["ok"] is True and body["findings"] == []


def test_object_info_cached_within_ttl(client, endpoint_at, fake_comfy):
    fake_comfy.object_info = build_object_info(_load_05())
    for _ in range(2):
        r = client.post(
            "/api/workflows/analyze",
            json={"workflow_json": _load_05(), "endpoint_id": endpoint_at},
        )
        assert r.json()["connection_check"]["ok"] is True
    hits = [p for p in fake_comfy.requests if p == "/object_info"]
    assert len(hits) == 1  # second call served from the 5-minute cache

    wf_id = _create_workflow(client, _load_05())
    client.post(f"/api/workflows/{wf_id}/check", json={"refresh": True})
    hits = [p for p in fake_comfy.requests if p == "/object_info"]
    assert len(hits) == 2  # manual refresh bypasses the TTL