Song367-ComfyUI-KieAPI/tests/test_client.py

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")