| import importlib.util |
| import json |
| import sys |
| import types |
| from pathlib import Path |
|
|
| import pytest |
|
|
|
|
| class _Port: |
| def __init__(self, name=None, **kwargs): |
| self.name = name |
| self.kwargs = kwargs |
|
|
|
|
| class _PortType: |
| @staticmethod |
| def Input(name, **kwargs): |
| return _Port(name, **kwargs) |
|
|
| @staticmethod |
| def Output(name=None, **kwargs): |
| return _Port(name, **kwargs) |
|
|
|
|
| class _NodeOutput: |
| def __init__(self, *values): |
| self.values = values |
|
|
|
|
| class _Schema: |
| def __init__(self, **kwargs): |
| self.__dict__.update(kwargs) |
|
|
|
|
| def _load_common_module(monkeypatch): |
| io = types.SimpleNamespace( |
| AnyType=_PortType, |
| Clip=_PortType, |
| ComfyNode=object, |
| Custom=lambda **kwargs: _PortType, |
| Hidden=types.SimpleNamespace(extra_pnginfo="EXTRA_PNGINFO"), |
| Int=_PortType, |
| Model=_PortType, |
| NodeOutput=_NodeOutput, |
| Schema=_Schema, |
| String=_PortType, |
| Vae=_PortType, |
| ) |
| comfy_api = types.ModuleType("comfy_api") |
| comfy_api_latest = types.ModuleType("comfy_api.latest") |
| comfy_api_latest.io = io |
| comfy_api.latest = comfy_api_latest |
| monkeypatch.setitem(sys.modules, "comfy_api", comfy_api) |
| monkeypatch.setitem(sys.modules, "comfy_api.latest", comfy_api_latest) |
|
|
| path = Path(__file__).parents[1] / "nodes" / "common.py" |
| spec = importlib.util.spec_from_file_location("easy_media_common_test", path) |
| assert spec is not None and spec.loader is not None |
| module = importlib.util.module_from_spec(spec) |
| spec.loader.exec_module(module) |
| return module |
|
|
|
|
| def test_model_loader_pack_schema_uses_fastuse_compatible_type(monkeypatch): |
| module = _load_common_module(monkeypatch) |
|
|
| schema = module.EasyModelLoaderPack.define_schema() |
| inputs = {port.name: port for port in schema.inputs} |
|
|
| assert schema.node_id == "easy modelLoaderPack" |
| assert list(inputs) == [ |
| "model", |
| "clip", |
| "vae", |
| "audio_vae", |
| ] |
| assert inputs["audio_vae"].kwargs["optional"] is True |
| assert schema.outputs[0].name == "model_loader" |
|
|
|
|
| def test_model_loader_pack_builds_expected_dictionary(monkeypatch): |
| module = _load_common_module(monkeypatch) |
| model = object() |
| clip = object() |
| vae = object() |
| audio_vae = object() |
| result = module.EasyModelLoaderPack.execute( |
| model, |
| clip, |
| vae, |
| audio_vae, |
| ) |
|
|
| assert result.values == ( |
| { |
| "model": model, |
| "clip": clip, |
| "vae": vae, |
| "audio_vae": audio_vae, |
| }, |
| ) |
|
|
|
|
| def test_model_loader_pack_omits_unconnected_optional_values(monkeypatch): |
| module = _load_common_module(monkeypatch) |
|
|
| model_loader = module.EasyModelLoaderPack.execute( |
| object(), object(), object() |
| ).values[0] |
|
|
| assert set(model_loader) == {"model", "clip", "vae"} |
|
|
|
|
| def test_model_loader_pack_includes_audio_vae_when_connected(monkeypatch): |
| module = _load_common_module(monkeypatch) |
| audio_vae = object() |
|
|
| model_loader = module.EasyModelLoaderPack.execute( |
| object(), object(), object(), audio_vae=audio_vae |
| ).values[0] |
|
|
| assert model_loader["audio_vae"] is audio_vae |
|
|
|
|
| def test_model_loader_pack_has_matching_chinese_localization(): |
| locale_path = Path(__file__).parents[1] / "locales" / "zh" / "nodeDefs.json" |
| if not locale_path.exists(): |
| pytest.skip("locales/zh/nodeDefs.json is not bundled in this repository") |
| node_defs = json.loads(locale_path.read_text()) |
| translation = node_defs["easy modelLoaderPack"] |
|
|
| assert translation["display_name"] == "模型加载器打包" |
| assert set(translation["inputs"]) == { |
| "model", |
| "clip", |
| "vae", |
| "audio_vae", |
| } |
| assert translation["outputs"] == {"0": {"name": "模型加载器"}} |
|
|
