xiaoxia-saas/tests/render_compare/create_test_asset.py

#!/usr/bin/env python3
"""
自动化测试用：一键创建 ready 状态的视频素材。

跳过转码和上传流程，直接指定 storage_key 创建可用素材，
供渲染对比测试等场景快速生成测试素材。

用法示例：
    python create_test_asset.py \
        --base-url http://localhost:8000 \
        --token xxx \
        --project-id proj_xxx \
        --name "测试素材-30s" \
        --storage-key "assets/test/video.mp4" \
        --duration 30 \
        --file-size 10485760
"""

import argparse
import json
import os
import sys
import urllib.error
import urllib.request


def parse_args():
    parser = argparse.ArgumentParser(description="创建 ready 状态的测试素材（跳过转码上传）")
    parser.add_argument(
        "--base-url",
        default=os.environ.get("API_BASE_URL", "http://localhost:8000"),
        help="API 地址，默认 http://localhost:8000（或环境变量 API_BASE_URL）",
    )
    parser.add_argument(
        "--token",
        default=os.environ.get("API_TOKEN", ""),
        help="登录 token（或环境变量 API_TOKEN）",
    )
    parser.add_argument(
        "--project-id",
        required=True,
        help="项目 ID",
    )
    parser.add_argument(
        "--name",
        required=True,
        help="素材名称",
    )
    parser.add_argument(
        "--storage-key",
        required=True,
        help="OSS storage_key（文件必须已存在于 OSS）",
    )
    parser.add_argument(
        "--duration",
        type=float,
        default=30.0,
        help="视频时长（秒），默认 30",
    )
    parser.add_argument(
        "--width",
        type=int,
        default=1280,
        help="视频宽度，默认 1280",
    )
    parser.add_argument(
        "--height",
        type=int,
        default=720,
        help="视频高度，默认 720",
    )
    parser.add_argument(
        "--fps",
        type=float,
        default=25.0,
        help="帧率，默认 25",
    )
    parser.add_argument(
        "--mime-type",
        default="video/mp4",
        help="MIME 类型，默认 video/mp4",
    )
    parser.add_argument(
        "--file-size",
        type=int,
        default=0,
        help="文件大小（字节），默认 0",
    )
    parser.add_argument(
        "--codec",
        default=None,
        help="视频编码，可选",
    )
    parser.add_argument(
        "--kind",
        default="video",
        choices=["video", "voice", "image"],
        help="素材库类型，默认 video",
    )
    return parser.parse_args()


def api_request(base_url: str, token: str, method: str, path: str, body: dict | None = None) -> dict:
    """发送 API 请求，返回 JSON 响应。"""
    url = f"{base_url.rstrip('/')}{path}"
    data = json.dumps(body).encode("utf-8") if body else None

    headers = {
        "Content-Type": "application/json",
    }
    if token:
        headers["Authorization"] = f"Bearer {token}"

    req = urllib.request.Request(url, data=data, method=method, headers=headers)

    try:
        with urllib.request.urlopen(req) as resp:
            resp_body = resp.read().decode("utf-8")
            return json.loads(resp_body) if resp_body else {}
    except urllib.error.HTTPError as e:
        error_body = e.read().decode("utf-8", errors="replace")
        print(f"[ERROR] API 请求失败: {method} {url}", file=sys.stderr)
        print(f"  HTTP {e.code}: {e.reason}", file=sys.stderr)
        print(f"  响应内容: {error_body}", file=sys.stderr)
        sys.exit(1)
    except urllib.error.URLError as e:
        print(f"[ERROR] 网络错误: {method} {url}", file=sys.stderr)
        print(f"  原因: {e.reason}", file=sys.stderr)
        sys.exit(1)


def ensure_default_library(base_url: str, token: str, project_id: str, kind: str) -> str:
    """确保项目有默认素材库，返回 library_id。"""
    print(f"[1/2] 确保默认 {kind} 素材库存在...")
    result = api_request(
        base_url,
        token,
        "POST",
        "/api/v1/asset-libraries/ensure-default",
        body={"project_id": project_id, "kind": kind},
    )
    library_id = result.get("id")
    library_name = result.get("name")
    print(f"      素材库: {library_name} (id: {library_id})")
    return library_id


def create_asset(
    base_url: str,
    token: str,
    project_id: str,
    library_id: str,
    name: str,
    storage_key: str,
    mime_type: str,
    file_size: int,
    duration: float,
    width: int,
    height: int,
    fps: float,
    codec: str | None,
) -> dict:
    """创建 ready 状态的素材。"""
    print("[2/2] 创建 ready 状态素材...")
    body = {
        "project_id": project_id,
        "library_id": library_id,
        "name": name,
        "storage_key": storage_key,
        "mime_type": mime_type,
        "file_size": file_size,
        "duration": duration,
        "width": width,
        "height": height,
        "fps": fps,
        "status": "ready",
        "classification_status": "pending",
        "metadata": {},
    }
    if codec:
        body["codec"] = codec

    result = api_request(
        base_url,
        token,
        "POST",
        "/api/v1/assets",
        body=body,
    )
    return result


def main():
    args = parse_args()

    if not args.token:
        print("[ERROR] 请通过 --token 参数或 API_TOKEN 环境变量提供登录 token", file=sys.stderr)
        sys.exit(1)

    print("=" * 60)
    print("创建测试素材工具")
    print("=" * 60)
    print(f"  API 地址: {args.base_url}")
    print(f"  项目 ID:  {args.project_id}")
    print(f"  素材名称: {args.name}")
    print(f"  storage_key: {args.storage_key}")
    print(f"  分辨率:   {args.width}x{args.height} @ {args.fps}fps")
    print(f"  时长:     {args.duration}s")
    print(f"  文件大小: {args.file_size} bytes")
    print("=" * 60)
    print()

    # Step 1: 确保默认素材库存在
    library_id = ensure_default_library(args.base_url, args.token, args.project_id, args.kind)

    # Step 2: 创建素材
    asset = create_asset(
        base_url=args.base_url,
        token=args.token,
        project_id=args.project_id,
        library_id=library_id,
        name=args.name,
        storage_key=args.storage_key,
        mime_type=args.mime_type,
        file_size=args.file_size,
        duration=args.duration,
        width=args.width,
        height=args.height,
        fps=args.fps,
        codec=args.codec,
    )

    asset_id = asset.get("id")
    print()
    print("=" * 60)
    print("✓ 素材创建成功！")
    print("=" * 60)
    print(f"  asset_id:   {asset_id}")
    print(f"  状态:       {asset.get('status')}")
    print(f"  素材库 ID:  {asset.get('library_id')}")
    print(f"  文件 URL:   {asset.get('file_url', 'N/A')}")
    print("=" * 60)
    print()
    print("调用示例：")
    print(f"  export ASSET_ID={asset_id}")
    print("  # 在生成任务中使用:")
    print("  #   --asset-id $ASSET_ID")
    print()

    # 输出 asset_id 到 stdout（方便脚本捕获）
    print(asset_id)


if __name__ == "__main__":
    main()