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", ]