test/test_gpt_psd.py

from __future__ import annotations

import json
import time
import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path
from uuid import uuid4

BASE_URL = "http://127.0.0.1:8000"
PROMPT = "按原图位置拆分海报元素并合成可编辑 PSD，同时输出每个图层素材 zip。"
BASE64_IMAGES: list[str] = []
TIMEOUT_SECS = 600
POLL_INTERVAL_SECS = 5


def request_json(method: str, path: str, payload: dict | None = None) -> dict:
    api_key = json.loads((Path(__file__).resolve().parents[1] / "config.json").read_text(encoding="utf-8"))["auth-key"]
    if not api_key.strip():
        raise ValueError("API_KEY is empty")
    data = None if payload is None else json.dumps(payload, ensure_ascii=False).encode()
    request = urllib.request.Request(
        BASE_URL.rstrip("/") + path,
        data=data,
        method=method,
        headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
    )
    try:
        with urllib.request.urlopen(request, timeout=60) as response:
            return json.loads(response.read().decode())
    except urllib.error.HTTPError as exc:
        raise RuntimeError(exc.read().decode("utf-8", "replace")) from exc


def main() -> None:
    task = request_json("POST", "/v1/psd/generations", {
        "client_task_id": str(uuid4()),
        "prompt": PROMPT,
        "base64_images": BASE64_IMAGES,
    })
    task_id = str(task.get("taskId") or task.get("id") or "")
    if not task_id:
        raise RuntimeError(f"missing taskId: {task}")
    print(json.dumps(task, ensure_ascii=False, indent=2))
    deadline = time.time() + TIMEOUT_SECS
    while time.time() < deadline:
        time.sleep(POLL_INTERVAL_SECS)
        status = request_json("GET", "/v1/editable-file-tasks?ids=" + urllib.parse.quote(task_id))
        print(json.dumps(status, ensure_ascii=False, indent=2))
        item = (status.get("items") or [{}])[0]
        if item.get("status") in {"success", "error"}:
            return
    raise TimeoutError(f"task timeout: {task_id}")


if __name__ == "__main__":
    main()