Tool_Agent/tests/test_kuaishou_kling.py

"""测试快手可灵AI工具

用法：
    uv run python tests/test_kuaishou_kling.py                    # Health check
    uv run python tests/test_kuaishou_kling.py --text2image       # 测试文生图
    uv run python tests/test_kuaishou_kling.py --text2video       # 测试文生视频
    uv run python tests/test_kuaishou_kling.py --lip-sync         # 测试对口型
"""

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 ===")
    print("  Calling kuaishou_kling...")
    try:
        resp = httpx.post(f"{BASE_URL}/select_tool", json={
            "tool_id": "kuaishou_kling",
            "params": {
                "biz_type": "aiImage",
                "prompt": "cute cat playing in garden",
                "aspect_ratio": "16:9",
                "image_count": 2
            }
        }, 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("  [PASS]")
        else:
            print(f"  ERROR : {data.get('error')}")
            print("  [FAIL]")
    except Exception as e:
        print(f"  ERROR : {e}")
        print("  [FAIL]")


def test_text2video():
    print("=== Test Text to Video ===")
    print("  Calling kuaishou_kling...")
    try:
        resp = httpx.post(f"{BASE_URL}/select_tool", json={
            "tool_id": "kuaishou_kling",
            "params": {
                "biz_type": "aiVideo",
                "prompt": "ocean waves at sunset",
                "aspect_ratio": "16:9"
            }
        }, timeout=300)
        print(f"  Status : {resp.status_code}")
        data = resp.json()
        if data["status"] == "success":
            print(f"  task_id: {data['result'].get('task_id')}")
            print("  [PASS]")
        else:
            print(f"  ERROR : {data.get('error')}")
            print("  [FAIL]")
    except Exception as e:
        print(f"  ERROR : {e}")
        print("  [FAIL]")


def test_lip_sync():
    print("=== Test AI Lip Sync ===")
    print("  Calling kuaishou_kling...")
    try:
        resp = httpx.post(f"{BASE_URL}/select_tool", json={
            "tool_id": "kuaishou_kling",
            "params": {
                "biz_type": "aiLipSync",
                "mode": "text2video",
                "text": "Hello world"
            }
        }, timeout=300)
        print(f"  Status : {resp.status_code}")
        data = resp.json()
        if data["status"] == "success":
            print("  [PASS]")
        else:
            print(f"  ERROR : {data.get('error')}")
            print("  [FAIL]")
    except Exception as e:
        print(f"  ERROR : {e}")
        print("  [FAIL]")


def main():
    parser = argparse.ArgumentParser(description="Kuaishou Kling AI Tool Test")
    parser.add_argument("--text2image", action="store_true")
    parser.add_argument("--text2video", action="store_true")
    parser.add_argument("--lip-sync", action="store_true")
    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.text2video:
        print()
        test_text2video()
        ran_any = True
    if args.lip_sync:
        print()
        test_lip_sync()
        ran_any = True

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

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


if __name__ == "__main__":
    main()