vantang-ComfyUI-MediaHub/scripts/smoke_comfyui.py

#!/usr/bin/env python3
"""Validate MediaHub's live ComfyUI node schema and API workflow alignment."""

import argparse
import json
import time
from pathlib import Path
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen


ROOT = Path(__file__).resolve().parents[1]
WORKFLOWS = sorted((ROOT / "tests" / "fixtures").glob("*_prompt.json"))
API_ENV = ROOT / ".api_env"


def request_json(url: str, payload: dict | None = None) -> dict:
    data = None if payload is None else json.dumps(payload).encode()
    request = Request(
        url,
        data=data,
        headers={"Content-Type": "application/json"},
        method="GET" if payload is None else "POST",
    )
    with urlopen(request, timeout=10) as response:
        return json.load(response)


def load_api_env(path: Path) -> dict[str, str]:
    values = {}
    if not path.exists():
        return values
    for line_number, raw_line in enumerate(path.read_text().splitlines(), start=1):
        line = raw_line.strip()
        if not line or line.startswith("#"):
            continue
        if "=" not in line:
            raise RuntimeError(f"Invalid {path.name} line {line_number}")
        name, value = line.split("=", 1)
        values[name.strip()] = value.strip().strip("'\"")
    return values


def wait_for_prompt(base_url: str, prompt_id: str, timeout: int = 180) -> dict:
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        history = request_json(f"{base_url}/history/{prompt_id}")
        result = history.get(prompt_id)
        if result:
            return result
        time.sleep(1)
    raise RuntimeError(f"Timed out waiting for ComfyUI prompt {prompt_id}")


def validate_workflow(
    base_url: str,
    workflow_path: Path,
    node_classes: list[str],
    api_key: str = "",
) -> dict:
    """Validate that a fixture's node inputs match the live ComfyUI schema."""
    workflow = json.loads(workflow_path.read_text())
    if api_key:
        for node in workflow.values():
            if node["class_type"] == "MediaHubAPIKey":
                node["inputs"]["api_key"] = api_key

    for node_class in node_classes:
        object_info = request_json(f"{base_url}/object_info/{node_class}")
        live_inputs = object_info[node_class]["input"]
        live_required = set(live_inputs["required"])
        live_accepted = live_required | set(live_inputs.get("optional") or {})

        target = next(
            node for node in workflow.values() if node["class_type"] == node_class
        )
        example_required = set(target["inputs"])
        if not live_required <= example_required or not example_required <= live_accepted:
            missing = sorted(live_required - example_required)
            extra = sorted(example_required - live_accepted)
            raise RuntimeError(
                f"{node_class} API workflow/schema mismatch: "
                f"missing={missing}, extra={extra}"
            )

    result = request_json(f"{base_url}/prompt", {"prompt": workflow})
    node_errors = result.get("node_errors") or {}
    if node_errors:
        raise RuntimeError(
            f"ComfyUI rejected {workflow_path.name}:\n"
            + json.dumps(node_errors, ensure_ascii=False, indent=2)
        )

    return result


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--base-url", default="http://127.0.0.1:8000")
    parser.add_argument(
        "--live",
        action="store_true",
        help="Use MINIMAX_CN_API_KEY from .api_env and wait for a real result.",
    )
    args = parser.parse_args()
    base_url = args.base_url.rstrip("/")

    try:
        env = {}
        if args.live:
            env = load_api_env(API_ENV)

        results = []
        for workflow_path in WORKFLOWS:
            workflow = json.loads(workflow_path.read_text())
            node_classes = sorted(
                {
                    node["class_type"]
                    for node in workflow.values()
                    if node["class_type"].startswith(("MediaHub", "MiniMax", "Agnes"))
                }
            )
            api_key = ""
            if args.live:
                key_name = (
                    "AGNES_API_KEY"
                    if workflow_path.name.startswith("agnes_")
                    else "MINIMAX_CN_API_KEY"
                )
                api_key = env.get(key_name, "")
                if not api_key:
                    raise RuntimeError(f"{key_name} is empty or missing in .api_env")
            results.append(
                (workflow_path.stem, validate_workflow(base_url, workflow_path, node_classes, api_key))
            )

        if args.live:
            for label, result in results:
                prompt_id = result.get("prompt_id")
                if not prompt_id:
                    raise RuntimeError(f"ComfyUI did not return a prompt_id for {label}")
                history = wait_for_prompt(base_url, prompt_id)
                status = history.get("status") or {}
                if status.get("status_str") != "success":
                    messages = status.get("messages") or []
                    raise RuntimeError(
                        f"Live {label} prompt failed:\n"
                        + json.dumps(messages, ensure_ascii=False, indent=2)
                    )
                outputs = history.get("outputs") or {}
                if not outputs:
                    raise RuntimeError(f"Live {label} prompt completed without outputs")

    except (HTTPError, URLError, KeyError, RuntimeError) as exc:
        print(f"FAIL: {exc}")
        return 1

    print(
        "PASS: live node schemas match API examples and ComfyUI accepted "
        "all prompt structures."
    )
    if args.live:
        print(
            "PASS: live fixture requests completed successfully."
        )
    else:
        print(
            "Queued jobs are expected to stop at an empty provider API key; "
            "no paid API request is made."
        )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())