Tool_Agent/tests/test_router_api.py

"""测试 Router 核心接口

用法：
    uv run python tests/test_router_api.py                               # 只跑 health check
    uv run python tests/test_router_api.py --search                      # 搜索工具列表
    uv run python tests/test_router_api.py --search image                # 关键词搜索
    uv run python tests/test_router_api.py --status                      # 工具运行状态
    uv run python tests/test_router_api.py --select image_stitcher       # POST /run_tool 调用工具
    uv run python tests/test_router_api.py --stitch                      # 测试图片拼接
    uv run python tests/test_router_api.py --create                      # 默认任务
    uv run python tests/test_router_api.py --create image_stitcher       # 指定任务文件
    uv run python tests/test_router_api.py --launch-env                  # 创建 RunComfy 启动环境工具
    uv run python tests/test_router_api.py --run-only                    # 创建 RunComfy 任务执行工具
    uv run python tests/test_router_api.py --stop-env                    # 创建 RunComfy 环境销毁工具
"""

import argparse
import base64
import json
import sys
import time
from pathlib import Path
from typing import Any

import httpx

BASE_URL = "http://127.0.0.1:8001"
TASKS_DIR = Path(__file__).parent / "tasks"
TEST_IMAGES_DIR = TASKS_DIR / "stitcher_images"
OUTPUT_DIR = Path(__file__).parent / "output"


def check_connection():
    try:
        httpx.get(f"{BASE_URL}/health", timeout=3)
    except httpx.ConnectError:
        print(f"ERROR: Cannot connect to {BASE_URL}")
        print("Please start the service first:")
        print("  uv run python -m tool_agent")
        sys.exit(1)


def test_health():
    print("=== Health Check ===")
    resp = httpx.get(f"{BASE_URL}/health")
    print(f"  Status : {resp.status_code}")
    print(f"  Body   : {json.dumps(resp.json(), ensure_ascii=False, indent=4)}")
    assert resp.status_code == 200
    print("  [PASS]")


def test_search_tools(keyword: str = None):
    print(f"=== Search Tools{f' (keyword={keyword!r})' if keyword else ''} ===")
    payload = {"keyword": keyword} if keyword else {}
    resp = httpx.post(f"{BASE_URL}/search_tools", json=payload)
    print(f"  Status : {resp.status_code}")
    if resp.status_code != 200:
        print(f"  Body   : {resp.text}")
        print("  [FAIL]")
        return
    data = resp.json()
    print(f"  Total  : {data['total']}")
    for t in data["tools"]:
        print(f"\n  [{t['tool_id']}]")
        print(f"    name          : {t['name']}")
        print(f"    category      : {t.get('category', '')}")
        print(f"    state         : {t['state']}")
        print(f"    runtime       : {t.get('runtime_type', '')}  host_dir={t.get('host_dir', '')}")
        print(f"    endpoint      : {t.get('http_method', '')} {t.get('endpoint_path', '')}  port={t.get('port')}")
        print(f"    stream_support: {t.get('stream_support', False)}")
        print(f"    description   : {t.get('description', '')}")
        print(f"    params ({len(t.get('params', []))}):")
        for p in t.get("params", []):
            req_mark = "*" if p["required"] else " "
            default_str = f"  default={p['default']}" if p.get("default") is not None else ""
            enum_str = f"  enum={p['enum']}" if p.get("enum") else ""
            print(f"      {req_mark} {p['name']:<25} {p['type']:<12} {p.get('description', '')}{default_str}{enum_str}")
        if t.get("output_schema"):
            out_props = t["output_schema"].get("properties", {})
            print(f"    output ({len(out_props)}):")
            for oname, odef in out_props.items():
                print(f"        {oname:<25} {odef.get('type', ''):<12} {odef.get('description', '')}")
    print("\n  [PASS]")


def test_tools_status():
    print("=== Tools Status ===")
    resp = httpx.get(f"{BASE_URL}/tools/status")
    print(f"  Status : {resp.status_code}")
    data = resp.json()
    print(f"  Total  : {len(data['tools'])}")
    for t in data["tools"]:
        print(f"    - {t['tool_id']}")
        print(f"        state  : {t['state']}")
        print(f"        port   : {t.get('port')}")
        print(f"        pid    : {t.get('pid')}")
        print(f"        sources: {[s['type'] for s in t.get('sources', [])]}")
        if t.get("last_error"):
            print(f"        error  : {t['last_error']}")
    print("  [PASS]")


def _run_tool(
    tool_id: str, params: dict[str, Any], timeout: float = 120.0
) -> tuple[bool, str | None, Any]:
    """POST /run_tool。成功返回 (True, None, result)；失败 (False, message, None)。"""
    resp = httpx.post(
        f"{BASE_URL}/run_tool",
        json={"tool_id": tool_id, "params": params},
        timeout=timeout,
    )
    print(f"  Status : {resp.status_code}")
    if resp.status_code != 200:
        return False, f"HTTP {resp.status_code}: {resp.text[:300]}", None
    try:
        data = resp.json()
    except Exception as e:
        return False, f"Invalid JSON: {e}", None
    if data.get("status") != "success":
        return False, data.get("error") or str(data), None
    result = data.get("result")
    if isinstance(result, dict) and result.get("status") == "error":
        return False, str(result.get("error", result)), None
    return True, None, result


def test_select_tool(tool_id: str):
    print(f"=== Run Tool (tool_id={tool_id!r}) ===")
    ok, err, result = _run_tool(tool_id, {}, timeout=30)
    print(f"  Result :")
    if not ok:
        print(f"    error : {err}")
        print("  [FAIL]")
        return
    result_str = json.dumps(result, ensure_ascii=False, indent=6)
    print(f"    body: {result_str[:500]}")
    print("  [PASS]")


def test_stitch_images():
    print("=== Test Image Stitcher ===")
    if not TEST_IMAGES_DIR.exists():
        print(f"  ERROR: Test images directory not found: {TEST_IMAGES_DIR}")
        print("  [SKIP]")
        return

    image_files = sorted(TEST_IMAGES_DIR.glob("*.png"))
    if len(image_files) < 2:
        print(f"  ERROR: Need at least 2 images, found {len(image_files)}")
        print("  [SKIP]")
        return

    print(f"  Images : {len(image_files)} found")
    images_b64 = []
    for img_path in image_files[:6]:
        with open(img_path, "rb") as f:
            images_b64.append(base64.b64encode(f.read()).decode())
        print(f"    - {img_path.name}")

    print(f"  Calling image_stitcher (grid, 2 columns)...")
    try:
        ok, err, result = _run_tool(
            "image_stitcher",
            {
                "images": images_b64,
                "direction": "grid",
                "columns": 2,
                "spacing": 10,
                "background_color": "#FFFFFF",
            },
            timeout=120.0,
        )
        if not ok:
            print(f"  ERROR : {err}")
            print("  [FAIL]")
            return
        if not isinstance(result, dict) or "image" not in result:
            print(f"  ERROR : 缺少 image 字段: {result!r}")
            print("  [FAIL]")
            return
        print(f"  Result :")
        print(f"    width : {result.get('width')}")
        print(f"    height: {result.get('height')}")
        OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
        output_path = OUTPUT_DIR / "stitched_result.png"
        with open(output_path, "wb") as f:
            f.write(base64.b64decode(result["image"]))
        print(f"    saved : {output_path}")
        print("  [PASS]")
    except httpx.TimeoutException:
        print("  ERROR : Request timeout")
        print("  [FAIL]")
    except Exception as e:
        print(f"  ERROR : {e}")
        print("  [FAIL]")


def load_task_spec(task_name: str) -> dict:
    task_file = TASKS_DIR / f"{task_name}.json"
    if not task_file.exists():
        print(f"  ERROR: Task file not found: {task_file}")
        print("  Available tasks:")
        if TASKS_DIR.exists():
            for f in TASKS_DIR.glob("*.json"):
                print(f"    - {f.stem}")
        sys.exit(1)
    with open(task_file, "r", encoding="utf-8") as f:
        return json.load(f)


def test_create_tool(task_name: str = None):
    print(f"=== Create Tool{f' (task={task_name!r})' if task_name else ''} ===")
    if task_name:
        task_data = load_task_spec(task_name)
        print(f"  File       : tests/tasks/{task_name}.json")
        print(f"  Description: {task_data['description'][:80]}")
    else:
        task_data = {"description": "创建一个简单的文本计数工具，输入文本，返回字数和字符数"}
        print(f"  Description: {task_data['description']}")

    resp = httpx.post(f"{BASE_URL}/create_tool", json=task_data)
    data = resp.json()
    task_id = data["task_id"]
    print(f"  Task ID    : {task_id}")
    print(f"  Status     : {data['status']}")
    assert data["status"] == "pending"
    print("  [SUBMITTED]")

    print(f"\n  Polling task {task_id} (timeout 10min)...")
    for i in range(120):
        time.sleep(5)
        resp = httpx.get(f"{BASE_URL}/tasks/{task_id}", timeout=30)
        task = resp.json()
        status = task["status"]

        if i % 6 == 0:
            print(f"    [{i*5}s] status={status}")

        if status == "completed":
            print(f"\n  Completed!")
            print(f"  Result     : {str(task.get('result', ''))[:300]}")
            resp2 = httpx.post(f"{BASE_URL}/search_tools", json={})
            tools = resp2.json()["tools"]
            print(f"  Registered : {[t['tool_id'] for t in tools]}")
            print("  [PASS]")
            return

        if status == "failed":
            print(f"\n  Failed!")
            print(f"  Error      : {task.get('error', 'unknown')}")
            print("  [FAIL]")
            return

    print(f"\n  Timeout after 600s")
    print("  [TIMEOUT]")


def main():
    parser = argparse.ArgumentParser(description="Router API Test")
    parser.add_argument("--search", nargs="?", const="", metavar="KEYWORD",
                        help="search tools, optional keyword")
    parser.add_argument("--status", action="store_true",
                        help="show tools status")
    parser.add_argument("--select", metavar="TOOL_ID",
                        help="call a tool by tool_id")
    parser.add_argument("--stitch", action="store_true",
                        help="test image stitcher with sample images")
    parser.add_argument("--create", nargs="?", const="", metavar="TASK_NAME",
                        help="create tool, optional task file name")
    parser.add_argument("--launch-env", action="store_true",
                        help="create RunComfy launch env tool (runcomfy_launch_env)")
    parser.add_argument("--run-only", action="store_true",
                        help="create RunComfy run only tool (runcomfy_run_only)")
    parser.add_argument("--stop-env", action="store_true",
                        help="create RunComfy stop env tool (runcomfy_stop_env)")
    args = parser.parse_args()

    print(f"Target: {BASE_URL}\n")
    check_connection()

    # 始终跑 health check
    test_health()

    ran_any = False

    if args.search is not None:
        print()
        test_search_tools(args.search or None)
        ran_any = True

    if args.status:
        print()
        test_tools_status()
        ran_any = True

    if args.select:
        print()
        test_select_tool(args.select)
        ran_any = True

    if args.stitch:
        print()
        test_stitch_images()
        ran_any = True

    if args.create is not None:
        print()
        test_create_tool(args.create or None)
        ran_any = True

    if args.launch_env:
        print()
        test_create_tool("runcomfy_launch_env")
        ran_any = True

    if args.run_only:
        print()
        test_create_tool("runcomfy_run_only")
        ran_any = True

    if args.stop_env:
        print()
        test_create_tool("runcomfy_stop_env")
        ran_any = True

    if not ran_any:
        print()
        print("No test specified. Available options:")
        parser.print_help()

    print("\n=== DONE ===")


if __name__ == "__main__":
    main()