Tool_Agent/tests/test_midjourney.py

"""测试 Midjourney 代理工具 — Router POST /run_tool

本地服务将 JSON 原样转发至 MIDJOURNEY_API_BASE 上你已实现的三个接口：
  POST /submit_job       cookie, prompt, user_id, mode(relaxed|fast)
  POST /query_job_status cookie, job_id
  POST /get_image_urls   job_id → 四张图链接

用法：
    1. tools/local/midjourney/.env：MIDJOURNEY_API_BASE
    2. uv run python -m tool_agent
    3. uv run python tests/test_midjourney.py

端到端（可选）：设置 MIDJOURNEY_TEST_COOKIE、MIDJOURNEY_TEST_USER_ID 后脚本会
    submit → 轮询 query → get_image_urls；否则仅校验工具已注册并退出 0。

环境变量：
    TOOL_AGENT_ROUTER_URL
    MIDJOURNEY_SUBMIT_TOOL_ID      默认 midjourney_submit_job
    MIDJOURNEY_QUERY_TOOL_ID       默认 midjourney_query_job_status
    MIDJOURNEY_GET_URLS_TOOL_ID    默认 midjourney_get_image_urls
    MIDJOURNEY_TEST_COOKIE / MIDJOURNEY_TEST_USER_ID / MIDJOURNEY_TEST_PROMPT / MIDJOURNEY_TEST_MODE
    MIDJOURNEY_POLL_INTERVAL_S / MIDJOURNEY_POLL_MAX_WAIT_S
"""

from __future__ import annotations

import io
import os
import sys
import time
from typing import Any

if sys.platform == "win32":
    _out = sys.stdout
    if isinstance(_out, io.TextIOWrapper):
        _out.reconfigure(encoding="utf-8")

import httpx

ROUTER_URL = os.environ.get("TOOL_AGENT_ROUTER_URL", "http://127.0.0.1:8001")
T_SUBMIT = os.environ.get("MIDJOURNEY_SUBMIT_TOOL_ID", "midjourney_submit_job")
T_QUERY = os.environ.get("MIDJOURNEY_QUERY_TOOL_ID", "midjourney_query_job_status")
T_URLS = os.environ.get("MIDJOURNEY_GET_URLS_TOOL_ID", "midjourney_get_image_urls")
TEST_COOKIE = os.environ.get("MIDJOURNEY_TEST_COOKIE", "").strip()
TEST_USER_ID = os.environ.get("MIDJOURNEY_TEST_USER_ID", "").strip()
TEST_PROMPT = os.environ.get("MIDJOURNEY_TEST_PROMPT", "a red apple on white background --v 6")
TEST_MODE = os.environ.get("MIDJOURNEY_TEST_MODE", "fast").strip().lower()
POLL_INTERVAL_S = float(os.environ.get("MIDJOURNEY_POLL_INTERVAL_S", "3"))
POLL_MAX_WAIT_S = float(os.environ.get("MIDJOURNEY_POLL_MAX_WAIT_S", "600"))


def run_tool(tool_id: str, params: dict[str, Any], timeout: float = 120.0) -> Any:
    resp = httpx.post(
        f"{ROUTER_URL}/run_tool",
        json={"tool_id": tool_id, "params": params},
        timeout=timeout,
    )
    resp.raise_for_status()
    body = resp.json()
    if body.get("status") != "success":
        raise RuntimeError(body.get("error") or str(body))
    result = body.get("result")
    if isinstance(result, dict) and result.get("status") == "error":
        raise RuntimeError(result.get("error", str(result)))
    return result


def _extract_job_id(data: dict[str, Any]) -> str | None:
    if not isinstance(data, dict):
        return None
    for key in ("job_id", "jobId", "id", "task_id", "taskId"):
        v = data.get(key)
        if v is not None and str(v).strip():
            return str(v).strip()
    inner = data.get("data")
    if isinstance(inner, dict):
        return _extract_job_id(inner)
    return None


def _status_terminal_ok(data: dict[str, Any]) -> bool:
    if not isinstance(data, dict):
        return False
    s = str(
        data.get("status")
        or data.get("job_status")
        or data.get("jobStatus")
        or data.get("state")
        or ""
    ).lower()
    if not s and isinstance(data.get("data"), dict):
        return _status_terminal_ok(data["data"])
    return any(k in s for k in ("complete", "success", "done", "finished", "succeed", "ready"))


def _status_terminal_fail(data: dict[str, Any]) -> bool:
    if not isinstance(data, dict):
        return False
    s = str(data.get("status") or data.get("job_status") or data.get("state") or "").lower()
    return any(k in s for k in ("fail", "error", "cancel", "canceled", "cancelled"))


def _extract_url_list(payload: Any) -> list[str]:
    if isinstance(payload, list):
        return [str(x) for x in payload if isinstance(x, str) and x.startswith("http")]
    if not isinstance(payload, dict):
        return []
    for key in ("image_urls", "urls", "images", "data"):
        v = payload.get(key)
        if isinstance(v, list):
            out = [str(x) for x in v if isinstance(x, str) and x.startswith("http")]
            if out:
                return out
        if isinstance(v, dict):
            nested = _extract_url_list(v)
            if nested:
                return nested
    return _extract_url_list(payload.get("data"))


def main() -> None:
    print("=" * 50)
    print("测试 Midjourney（submit / query / get_image_urls）")
    print("=" * 50)
    print(f"ROUTER_URL: {ROUTER_URL}")

    try:
        r = httpx.get(f"{ROUTER_URL}/health", timeout=3)
        print(f"Router 状态: {r.json()}")
    except httpx.ConnectError:
        print(f"无法连接 Router ({ROUTER_URL})，请先: uv run python -m tool_agent")
        sys.exit(1)

    print("\n--- 校验工具已注册 ---")
    tr = httpx.get(f"{ROUTER_URL}/tools", timeout=30)
    tr.raise_for_status()
    tools = tr.json().get("tools", [])
    ids = {t["tool_id"] for t in tools}
    for tid in (T_SUBMIT, T_QUERY, T_URLS):
        if tid not in ids:
            print(f"错误: {tid!r} 不在 GET /tools 中。示例: {sorted(ids)[:25]}...")
            sys.exit(1)
        meta = next(t for t in tools if t["tool_id"] == tid)
        print(f"  {tid}: {meta.get('name', '')} (state={meta.get('state')})")

    if not TEST_COOKIE or not TEST_USER_ID:
        print(
            "\n未设置 MIDJOURNEY_TEST_COOKIE 与 MIDJOURNEY_TEST_USER_ID，跳过端到端；"
            "工具注册检查已通过，退出 0。"
        )
        return

    if TEST_MODE not in ("relaxed", "fast"):
        print(f"错误: MIDJOURNEY_TEST_MODE 须为 relaxed 或 fast，当前: {TEST_MODE!r}")
        sys.exit(1)

    print("\n--- midjourney_submit_job ---")
    try:
        sub = run_tool(
            T_SUBMIT,
            {
                "cookie": TEST_COOKIE,
                "prompt": TEST_PROMPT,
                "user_id": TEST_USER_ID,
                "mode": TEST_MODE,
            },
            timeout=180.0,
        )
    except (RuntimeError, httpx.HTTPError) as e:
        print(f"错误: {e}")
        sys.exit(1)

    if not isinstance(sub, dict):
        print(f"错误: submit 返回非 object: {type(sub)}")
        sys.exit(1)

    job_id = _extract_job_id(sub)
    if not job_id:
        print(f"错误: 无法从 submit 响应解析 job_id: {sub}")
        sys.exit(1)
    print(f"job_id: {job_id}")

    print("\n--- midjourney_query_job_status 轮询 ---")
    deadline = time.monotonic() + POLL_MAX_WAIT_S
    last: dict[str, Any] = {}

    while time.monotonic() < deadline:
        time.sleep(POLL_INTERVAL_S)
        try:
            q = run_tool(
                T_QUERY,
                {"cookie": TEST_COOKIE, "job_id": job_id},
                timeout=120.0,
            )
        except (RuntimeError, httpx.HTTPError) as e:
            print(f"轮询错误: {e}")
            sys.exit(1)

        last = q if isinstance(q, dict) else {}
        st = last.get("status") or last.get("job_status") or last.get("state")
        print(f"  status: {st}")

        if _status_terminal_fail(last):
            print(f"任务失败: {last}")
            sys.exit(1)
        if _status_terminal_ok(last):
            break
    else:
        print(f"等待超时 ({POLL_MAX_WAIT_S}s)，最后响应: {last}")
        sys.exit(1)

    print("\n--- midjourney_get_image_urls ---")
    try:
        urls_payload = run_tool(T_URLS, {"job_id": job_id}, timeout=120.0)
    except (RuntimeError, httpx.HTTPError) as e:
        print(f"错误: {e}")
        sys.exit(1)

    urls = _extract_url_list(urls_payload)
    if len(urls) < 4:
        print(f"警告: 期望至少 4 个 http 链接，实际 {len(urls)}；原始: {str(urls_payload)[:500]}")
        if len(urls) == 0:
            sys.exit(1)

    for i, u in enumerate(urls[:4], 1):
        print(f"  [{i}] {u[:96]}...")
    print("\n测试通过!")


if __name__ == "__main__":
    main()