195 lines
6.6 KiB
Python
195 lines
6.6 KiB
Python
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")
|