Tool_Agent/tests/test_jimeng.py

"""测试 ji_meng 任务工具 — 通过 Router API 调用（范本同 test_liblibai_tool.py）

用法：
    1. 先启动 Router：uv run python -m tool_agent
    2. 运行测试：python tests/test_jimeng.py

需已注册 ji_meng_add_task、ji_meng_query_task（或改下方常量）；搜索关键词默认可匹配二者。
"""

import sys
import time

if sys.platform == 'win32':
    sys.stdout.reconfigure(encoding='utf-8')

import httpx

ROUTER_URL = "http://127.0.0.1:8001"
POLL_INTERVAL_S = 2
POLL_MAX_WAIT_S = 300


def _extract_task_id(data: dict):
    for key in ("task_id", "taskId", "id", "job_id", "jobId"):
        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_task_id(inner)
    return None


def _terminal_success(data: dict) -> bool:
    status = (
        data.get("status")
        or data.get("task_status")
        or data.get("taskStatus")
        or data.get("state")
    )
    if status is None and isinstance(data.get("data"), dict):
        status = data["data"].get("status") or data["data"].get("task_status")
    if status is None:
        return False
    s = str(status).lower()
    return s in ("completed", "success", "done", "finished", "succeed", "complete")


def _terminal_failure(data: dict) -> bool:
    status = data.get("status") or data.get("task_status") or data.get("state")
    if status is None and isinstance(data.get("data"), dict):
        status = data["data"].get("status")
    if status is None:
        return False
    s = str(status).lower()
    return s in ("failed", "error", "cancelled", "canceled")


def main():
    print("=" * 50)
    print("测试 ji_meng 任务工具")
    print("=" * 50)

    # 1. 检查 Router 是否在线
    try:
        resp = httpx.get(f"{ROUTER_URL}/health", timeout=3)
        print(f"Router 状态: {resp.json()}")
    except httpx.ConnectError:
        print(f"无法连接 Router ({ROUTER_URL})")
        print("请先启动: uv run python -m tool_agent")
        sys.exit(1)

    # 2. 搜索工具，确认已注册
    print("\n--- 搜索工具 ---")
    resp = httpx.post(f"{ROUTER_URL}/search_tools", json={"keyword": "ji_meng"})
    tools = resp.json()
    print(f"找到 {tools['total']} 个工具")
    for t in tools["tools"]:
        print(f"  {t['tool_id']}: {t['name']} (state={t['state']})")

    # 3. 创建 ji_meng_add_task 工具
    print("\n--- 调用 ji_meng_add_task 工具创建任务 ---")
    print(f"提示词: simple white line art, cat, black background")
    print("提交中...")

    resp = httpx.post(
        f"{ROUTER_URL}/select_tool",
        json={
            "tool_id": "ji_meng_add_task",
            "params": {
                "task_type": "image",
                "prompt": "simple white line art, cat, black background",
            },
        },
        timeout=120,
    )

    result = resp.json()
    print(f"\n响应状态: {result.get('status')}")

    if result.get("code") != 0:
        print(f"错误: {result.get('msg')}")
        sys.exit(1)

    task_id = result.get("data", {}).get("task_id")
    if not task_id:
        print("错误: 无法从创建任务响应中解析 task_id")
        sys.exit(1)
    print(f"任务 ID: {task_id}")

    # 4. 轮询查询任务
    print("\n--- 调用 ji_meng_query_task 工具查询任务 ---")
    deadline = time.monotonic() + POLL_MAX_WAIT_S
    last = {}

    while time.monotonic() < deadline:
        resp = httpx.post(
            f"{ROUTER_URL}/select_tool",
            json={
                "tool_id": "ji_meng_query_task",
                "params": {"task_id": task_id},
            },
            timeout=60,
        )
        result = resp.json()
        print(f"\n响应状态: {result.get('status')}")

        if result.get("status") != "success":
            print(f"错误: {result.get('error')}")
            sys.exit(1)

        last = result.get("result", {})
        if not isinstance(last, dict):
            print(f"非 dict 结果: {last}")
            time.sleep(POLL_INTERVAL_S)
            continue
        if last.get("status") == "error":
            print(f"错误: {last.get('error')}")
            sys.exit(1)

        print(f"任务状态: {last.get('status')}")
        print(f"最终结果: {last}")
        print("\n测试通过!")
        return

        time.sleep(POLL_INTERVAL_S)

    print(f"\n等待超时 ({POLL_MAX_WAIT_S}s)，最后一次响应: {last}")
    sys.exit(1)


if __name__ == "__main__":
    main()