ComfyUI/tests-unit/comfy_extras_test/test_seedvr2_nodes.py

import importlib
import inspect
import sys
from unittest.mock import MagicMock, patch

import torch

from comfy.cli_args import args as cli_args

if not torch.cuda.is_available():
    cli_args.cpu = True


def test_seedvr_node_signature_matches_schema():
    mock_mm = MagicMock()
    mock_mm.xformers_enabled.return_value = False
    mock_mm.xformers_enabled_vae.return_value = False
    mock_mm.sage_attention_enabled.return_value = False
    mock_mm.flash_attention_enabled.return_value = False

    sentinel = object()
    prior_cpu = cli_args.cpu
    cli_args.cpu = True
    prior_module = sys.modules.get("comfy_extras.nodes_seedvr", sentinel)
    comfy_pkg = sys.modules.get("comfy")
    prior_mm_attr = getattr(comfy_pkg, "model_management", sentinel) if comfy_pkg else sentinel

    with patch.dict(sys.modules, {"comfy.model_management": mock_mm}):
        if comfy_pkg is not None:
            setattr(comfy_pkg, "model_management", mock_mm)
        sys.modules.pop("comfy_extras.nodes_seedvr", None)
        try:
            nodes_seedvr = importlib.import_module("comfy_extras.nodes_seedvr")
            for node_cls in (nodes_seedvr.SeedVR2Preprocess, nodes_seedvr.SeedVR2PostProcessing, nodes_seedvr.SeedVR2Conditioning):
                schema_ids = [i.id for i in node_cls.define_schema().inputs]
                exec_params = [
                    p for p in inspect.signature(node_cls.execute).parameters.keys()
                    if p != "cls"
                ]
                assert schema_ids == exec_params, (
                    f"{node_cls.__name__} schema/execute drift: "
                    f"schema_ids={schema_ids}, exec_params={exec_params}"
                )
        finally:
            cli_args.cpu = prior_cpu
            if prior_module is sentinel:
                sys.modules.pop("comfy_extras.nodes_seedvr", None)
            else:
                sys.modules["comfy_extras.nodes_seedvr"] = prior_module
            if comfy_pkg is not None:
                if prior_mm_attr is sentinel:
                    if hasattr(comfy_pkg, "model_management"):
                        delattr(comfy_pkg, "model_management")
                else:
                    setattr(comfy_pkg, "model_management", prior_mm_attr)