from __future__ import annotations import asyncio import importlib.util import json from io import BytesIO from pathlib import Path import torch from PIL import Image PLUGIN_ROOT = Path(__file__).resolve().parents[1] PLUGIN_INIT = PLUGIN_ROOT / "__init__.py" def load_plugin_module(): assert PLUGIN_INIT.exists(), f"Expected plugin entrypoint at {PLUGIN_INIT}" spec = importlib.util.spec_from_file_location("comfyui_kieapi_plugin", PLUGIN_INIT) 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 load_support_module(name: str, relative_path: str): module_path = PLUGIN_ROOT / relative_path assert module_path.exists(), f"Expected module at {module_path}" spec = importlib.util.spec_from_file_location(name, module_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_plugin_exports_node_mappings(): module = load_plugin_module() assert hasattr(module, "NODE_CLASS_MAPPINGS") assert hasattr(module, "NODE_DISPLAY_NAME_MAPPINGS") assert "KieNanoBanana2" in module.NODE_CLASS_MAPPINGS assert module.NODE_DISPLAY_NAME_MAPPINGS["KieNanoBanana2"] == "Kie Nano Banana 2" def test_kie_nano_banana_2_node_schema(): module = load_plugin_module() node_class = module.NODE_CLASS_MAPPINGS["KieNanoBanana2"] input_types = node_class.INPUT_TYPES() assert "required" in input_types assert "optional" in input_types assert set(input_types["required"]).issuperset( { "prompt", "api_key", "aspect_ratio", "resolution", "output_format", "google_search", "poll_timeout_seconds", } ) assert "reference_image" in input_types["optional"] assert node_class.RETURN_TYPES == ("IMAGE", "STRING") assert node_class.RETURN_NAMES == ("images", "task_id") assert node_class.CATEGORY == "api/Kie" def test_kie_nano_banana_2_execute_returns_image_batch_and_task_id(monkeypatch): module = load_plugin_module() node_class = module.NODE_CLASS_MAPPINGS["KieNanoBanana2"] globals_dict = node_class.execute.__globals__ image = Image.new("RGB", (1, 1), color=(255, 0, 0)) buffer = BytesIO() image.save(buffer, format="PNG") raw_bytes = buffer.getvalue() class FakeClient: def __init__(self, api_key: str): self.api_key = api_key async def generate_images(self, **kwargs): assert self.api_key == "resolved-key" assert kwargs["prompt"] == "banana ship" return [raw_bytes], "task_123" monkeypatch.setitem(globals_dict, "KieClient", FakeClient) monkeypatch.setitem(globals_dict, "resolve_api_key", lambda value: "resolved-key") image_batch, task_id = asyncio.run( node_class().execute( prompt="banana ship", api_key="node-key", aspect_ratio="auto", resolution="2K", output_format="png", google_search=False, poll_timeout_seconds=600, ) ) assert task_id == "task_123" assert tuple(image_batch.shape) == (1, 1, 1, 3) assert torch.isclose(image_batch[0, 0, 0, 0], torch.tensor(1.0)) def test_kie_nano_banana_2_execute_forwards_all_batched_reference_images(monkeypatch): module = load_plugin_module() node_class = module.NODE_CLASS_MAPPINGS["KieNanoBanana2"] globals_dict = node_class.execute.__globals__ generated_image = Image.new("RGB", (1, 1), color=(0, 0, 255)) buffer = BytesIO() generated_image.save(buffer, format="PNG") raw_bytes = buffer.getvalue() reference_batch = torch.tensor( [ [[[1.0, 0.0, 0.0]]], [[[0.0, 1.0, 0.0]]], [[[0.0, 0.0, 1.0]]], ], dtype=torch.float32, ) class FakeClient: def __init__(self, api_key: str): self.api_key = api_key async def generate_images(self, **kwargs): assert self.api_key == "resolved-key" assert len(kwargs["reference_image_data_uris"]) == 3 assert all(uri.startswith("data:image/png;base64,") for uri in kwargs["reference_image_data_uris"]) return [raw_bytes], "task_batch_refs" monkeypatch.setitem(globals_dict, "KieClient", FakeClient) monkeypatch.setitem(globals_dict, "resolve_api_key", lambda value: "resolved-key") image_batch, task_id = asyncio.run( node_class().execute( prompt="banana ship", api_key="node-key", aspect_ratio="auto", resolution="2K", output_format="png", google_search=False, poll_timeout_seconds=600, reference_image=reference_batch, ) ) assert task_id == "task_batch_refs" assert tuple(image_batch.shape) == (1, 1, 1, 3) def test_resolve_api_key_prefers_node_input(monkeypatch, tmp_path): module = load_support_module("comfyui_kieapi_config", "lib/config.py") config_path = tmp_path / "config.json" config_path.write_text(json.dumps({"api_key": "config-key"}), encoding="utf-8") monkeypatch.setenv("KIE_API_KEY", "env-key") assert module.resolve_api_key("node-key", config_path=config_path) == "node-key" def test_resolve_api_key_prefers_env_over_config(monkeypatch, tmp_path): module = load_support_module("comfyui_kieapi_config", "lib/config.py") config_path = tmp_path / "config.json" config_path.write_text(json.dumps({"api_key": "config-key"}), encoding="utf-8") monkeypatch.setenv("KIE_API_KEY", "env-key") assert module.resolve_api_key("", config_path=config_path) == "env-key" def test_resolve_api_key_uses_config_file(monkeypatch, tmp_path): module = load_support_module("comfyui_kieapi_config", "lib/config.py") config_path = tmp_path / "config.json" config_path.write_text(json.dumps({"api_key": "config-key"}), encoding="utf-8") monkeypatch.delenv("KIE_API_KEY", raising=False) assert module.resolve_api_key("", config_path=config_path) == "config-key" def test_resolve_api_key_raises_when_missing(monkeypatch, tmp_path): module = load_support_module("comfyui_kieapi_config", "lib/config.py") config_path = tmp_path / "config.json" monkeypatch.delenv("KIE_API_KEY", raising=False) try: module.resolve_api_key("", config_path=config_path) except ValueError as exc: assert "api key" in str(exc).lower() else: raise AssertionError("Expected resolve_api_key to raise when no API key is available")