Tool_Agent/tests/test_jimeng_ai.py

"""测试即梦AI工具

用法：
    uv run python tests/test_jimeng_ai.py                    # Health check
    uv run python tests/test_jimeng_ai.py --text2image       # 测试文生图
    uv run python tests/test_jimeng_ai.py --image2video      # 测试图生视频
    uv run python tests/test_jimeng_ai.py --query            # 测试查询任务状态
"""

import argparse
import json
import sys
import httpx

BASE_URL = "http://127.0.0.1:8001"


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_text2image():
    print("=== Test Text to Image (Seendance 2.0) ===")
    print("  Calling jimeng_ai...")
    try:
        resp = httpx.post(f"{BASE_URL}/select_tool", json={
            "tool_id": "jimeng_ai",
            "params": {
                "action": "text2image",
                "prompt": "beautiful sunset over mountains, high quality",
                "model": "seendance_2.0",
                "aspect_ratio": "16:9",
                "image_count": 1,
                "cfg_scale": 7.0,
                "steps": 20
            }
        }, timeout=120)
        print(f"  Status : {resp.status_code}")
        data = resp.json()
        if data["status"] == "success":
            result = data["result"]
            print(f"  task_id: {result.get('task_id')}")
            print(f"  status : {result.get('status')}")
            if result.get("result", {}).get("images"):
                print(f"  images : {len(result['result']['images'])} generated")
            print("  [PASS]")
        else:
            print(f"  ERROR : {data.get('error')}")
            print("  [FAIL]")
    except Exception as e:
        print(f"  ERROR : {e}")
        print("  [FAIL]")


def test_image2video():
    print("=== Test Image to Video (Seedream Lite 5.0) ===")
    print("  Calling jimeng_ai...")
    try:
        resp = httpx.post(f"{BASE_URL}/select_tool", json={
            "tool_id": "jimeng_ai",
            "params": {
                "action": "image2video",
                "image_url": "https://example.com/sample.jpg",
                "prompt": "camera slowly zooming in",
                "model": "seedream_lite_5.0",
                "video_duration": 5,
                "motion_strength": 0.5
            }
        }, timeout=300)
        print(f"  Status : {resp.status_code}")
        data = resp.json()
        if data["status"] == "success":
            result = data["result"]
            print(f"  task_id: {result.get('task_id')}")
            print(f"  status : {result.get('status')}")
            if result.get("result", {}).get("videos"):
                print(f"  videos : {len(result['result']['videos'])} generated")
            print("  [PASS]")
        else:
            print(f"  ERROR : {data.get('error')}")
            print("  [FAIL]")
    except Exception as e:
        print(f"  ERROR : {e}")
        print("  [FAIL]")


def test_query_status():
    print("=== Test Query Task Status ===")
    print("  NOTE: Requires a valid task_id from previous generation")
    print("  [SKIP - Manual test required]")


def main():
    parser = argparse.ArgumentParser(description="Jimeng AI Tool Test")
    parser.add_argument("--text2image", action="store_true", help="test text to image")
    parser.add_argument("--image2video", action="store_true", help="test image to video")
    parser.add_argument("--query", action="store_true", help="test query status")
    args = parser.parse_args()

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

    ran_any = False

    if args.text2image:
        print()
        test_text2image()
        ran_any = True

    if args.image2video:
        print()
        test_image2video()
        ran_any = True

    if args.query:
        print()
        test_query_status()
        ran_any = True

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

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


if __name__ == "__main__":
    main()