261 lines
8.2 KiB
Python
261 lines
8.2 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import importlib.util
|
|
from pathlib import Path
|
|
|
|
|
|
PLUGIN_ROOT = Path(__file__).resolve().parents[1]
|
|
|
|
|
|
def load_client_module():
|
|
module_path = PLUGIN_ROOT / "lib" / "client.py"
|
|
assert module_path.exists(), f"Expected module at {module_path}"
|
|
spec = importlib.util.spec_from_file_location("comfyui_kieapi_client", 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_upload_reference_image_returns_download_url():
|
|
module = load_client_module()
|
|
client = module.KieClient(api_key="token-123")
|
|
captured = {}
|
|
|
|
async def fake_request(method, url, *, params=None, json_body=None):
|
|
captured["method"] = method
|
|
captured["url"] = url
|
|
captured["json_body"] = json_body
|
|
return {
|
|
"success": True,
|
|
"code": 200,
|
|
"msg": "File uploaded successfully",
|
|
"data": {
|
|
"fileName": "uploaded-image.png",
|
|
"filePath": "images/user-uploads/uploaded-image.png",
|
|
"downloadUrl": "https://tempfile.redpandaai.co/xxx/images/user-uploads/uploaded-image.png",
|
|
"fileSize": 154832,
|
|
"mimeType": "image/png",
|
|
"uploadedAt": "2025-01-01T12:00:00.000Z",
|
|
},
|
|
}
|
|
|
|
client._request_json = fake_request
|
|
|
|
download_url = asyncio.run(
|
|
client.upload_reference_image("data:image/png;base64,abc123", file_name="reference.png")
|
|
)
|
|
|
|
assert download_url == "https://tempfile.redpandaai.co/xxx/images/user-uploads/uploaded-image.png"
|
|
assert captured["method"] == "POST"
|
|
assert captured["url"].endswith("/api/file-base64-upload")
|
|
assert captured["json_body"]["fileName"] == "reference.png"
|
|
|
|
|
|
def test_create_task_returns_task_id():
|
|
module = load_client_module()
|
|
client = module.KieClient(api_key="token-123")
|
|
captured = {}
|
|
|
|
async def fake_request(method, url, *, params=None, json_body=None):
|
|
captured["method"] = method
|
|
captured["url"] = url
|
|
captured["json_body"] = json_body
|
|
return {
|
|
"code": 200,
|
|
"msg": "success",
|
|
"data": {
|
|
"taskId": "task_nano-banana-2_1765178625768",
|
|
},
|
|
}
|
|
|
|
client._request_json = fake_request
|
|
|
|
task_id = asyncio.run(
|
|
client.create_nano_banana_2_task(
|
|
prompt="banana ship",
|
|
aspect_ratio="auto",
|
|
resolution="2K",
|
|
output_format="jpg",
|
|
google_search=False,
|
|
image_input=["https://example.com/reference.png"],
|
|
)
|
|
)
|
|
|
|
assert task_id == "task_nano-banana-2_1765178625768"
|
|
assert captured["method"] == "POST"
|
|
assert captured["url"].endswith("/api/v1/jobs/createTask")
|
|
assert captured["json_body"]["model"] == "nano-banana-2"
|
|
assert captured["json_body"]["input"]["prompt"] == "banana ship"
|
|
|
|
|
|
def test_wait_for_task_transitions_from_waiting_to_success():
|
|
module = load_client_module()
|
|
client = module.KieClient(api_key="token-123")
|
|
responses = iter(
|
|
[
|
|
{
|
|
"code": 200,
|
|
"msg": "success",
|
|
"data": {
|
|
"taskId": "task_123",
|
|
"state": "waiting",
|
|
"resultJson": "",
|
|
"failCode": "",
|
|
"failMsg": "",
|
|
},
|
|
},
|
|
{
|
|
"code": 200,
|
|
"msg": "success",
|
|
"data": {
|
|
"taskId": "task_123",
|
|
"state": "success",
|
|
"resultJson": "{\"resultUrls\":[\"https://example.com/generated.jpg\"]}",
|
|
"failCode": "",
|
|
"failMsg": "",
|
|
},
|
|
},
|
|
]
|
|
)
|
|
sleep_calls = []
|
|
|
|
async def fake_request(method, url, *, params=None, json_body=None):
|
|
return next(responses)
|
|
|
|
async def fake_sleep(seconds):
|
|
sleep_calls.append(seconds)
|
|
|
|
client._request_json = fake_request
|
|
client._sleep = fake_sleep
|
|
|
|
result = asyncio.run(client.wait_for_task("task_123", timeout_seconds=10))
|
|
|
|
assert result.state == "success"
|
|
assert result.result_urls() == ["https://example.com/generated.jpg"]
|
|
assert sleep_calls == [2.0]
|
|
|
|
|
|
def test_wait_for_task_raises_failure_error():
|
|
module = load_client_module()
|
|
client = module.KieClient(api_key="token-123")
|
|
|
|
async def fake_request(method, url, *, params=None, json_body=None):
|
|
return {
|
|
"code": 200,
|
|
"msg": "success",
|
|
"data": {
|
|
"taskId": "task_123",
|
|
"state": "fail",
|
|
"resultJson": "",
|
|
"failCode": "content_blocked",
|
|
"failMsg": "prompt rejected",
|
|
},
|
|
}
|
|
|
|
client._request_json = fake_request
|
|
|
|
try:
|
|
asyncio.run(client.wait_for_task("task_123", timeout_seconds=10))
|
|
except RuntimeError as exc:
|
|
assert "content_blocked" in str(exc)
|
|
assert "prompt rejected" in str(exc)
|
|
else:
|
|
raise AssertionError("Expected wait_for_task to raise RuntimeError on failed tasks")
|
|
|
|
|
|
def test_wait_for_task_raises_timeout_error():
|
|
module = load_client_module()
|
|
client = module.KieClient(api_key="token-123")
|
|
now_values = iter([0.0, 0.0, 2.1, 2.1])
|
|
|
|
async def fake_request(method, url, *, params=None, json_body=None):
|
|
return {
|
|
"code": 200,
|
|
"msg": "success",
|
|
"data": {
|
|
"taskId": "task_123",
|
|
"state": "waiting",
|
|
"resultJson": "",
|
|
"failCode": "",
|
|
"failMsg": "",
|
|
},
|
|
}
|
|
|
|
async def fake_sleep(seconds):
|
|
return None
|
|
|
|
client._request_json = fake_request
|
|
client._sleep = fake_sleep
|
|
client._now = lambda: next(now_values)
|
|
|
|
try:
|
|
asyncio.run(client.wait_for_task("task_123", timeout_seconds=2))
|
|
except TimeoutError as exc:
|
|
assert "task_123" in str(exc)
|
|
else:
|
|
raise AssertionError("Expected wait_for_task to raise TimeoutError")
|
|
|
|
|
|
def test_generate_images_uploads_all_reference_images_in_order():
|
|
module = load_client_module()
|
|
client = module.KieClient(api_key="token-123")
|
|
uploaded = []
|
|
created_payload = {}
|
|
|
|
async def fake_upload(data_uri, *, file_name="reference.png"):
|
|
uploaded.append((data_uri, file_name))
|
|
return f"https://example.com/{file_name}"
|
|
|
|
async def fake_create_task(**kwargs):
|
|
created_payload.update(kwargs)
|
|
return "task_multi_ref"
|
|
|
|
async def fake_wait_for_task(task_id, *, timeout_seconds=600):
|
|
return module.KieTaskRecordData(
|
|
taskId=task_id,
|
|
state="success",
|
|
resultJson={"resultUrls": ["https://example.com/generated.png"]},
|
|
failCode="",
|
|
failMsg="",
|
|
)
|
|
|
|
async def fake_download_result_images(task):
|
|
return [b"image-bytes"]
|
|
|
|
client.upload_reference_image = fake_upload
|
|
client.create_nano_banana_2_task = fake_create_task
|
|
client.wait_for_task = fake_wait_for_task
|
|
client.download_result_images = fake_download_result_images
|
|
client.build_reference_file_name = lambda index: f"unique-reference-{index}.png"
|
|
|
|
raw_images, task_id = asyncio.run(
|
|
client.generate_images(
|
|
prompt="banana ship",
|
|
aspect_ratio="auto",
|
|
resolution="2K",
|
|
output_format="png",
|
|
google_search=False,
|
|
poll_timeout_seconds=600,
|
|
reference_image_data_uris=[
|
|
"data:image/png;base64,first",
|
|
"data:image/png;base64,second",
|
|
"data:image/png;base64,third",
|
|
],
|
|
)
|
|
)
|
|
|
|
assert task_id == "task_multi_ref"
|
|
assert raw_images == [b"image-bytes"]
|
|
assert uploaded == [
|
|
("data:image/png;base64,first", "unique-reference-0.png"),
|
|
("data:image/png;base64,second", "unique-reference-1.png"),
|
|
("data:image/png;base64,third", "unique-reference-2.png"),
|
|
]
|
|
assert created_payload["image_input"] == [
|
|
"https://example.com/unique-reference-0.png",
|
|
"https://example.com/unique-reference-1.png",
|
|
"https://example.com/unique-reference-2.png",
|
|
]
|