Tool_Agent/tools/local/veo_video_gen/tests/test_veo.py

import os
import sys
import requests
import json

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from veo_client import VEOClient

API_KEY = os.getenv("APIYI_KEY")
if not API_KEY:
    print("ERROR: APIYI_KEY not set")
    sys.exit(1)

os.makedirs(os.path.join(os.path.dirname(__file__), "output"), exist_ok=True)

print("=== VEO 3.1 API Test ===")
client = VEOClient(API_KEY)

# Step 1: create task
print("\n[1] Creating video task...")
video_id = client.create_video("A cat walking in a sunny garden", model="veo-3.1-fast")
print(f"    video_id: {video_id}")

# Step 2: poll until done
print("\n[2] Waiting for completion...")
result = client.wait_for_completion(video_id, timeout=600, interval=5)
video_url = result["video_url"]
print(f"    Done! URL: {video_url}")

# Step 3: download
print("\n[3] Downloading video...")
resp = requests.get(video_url, timeout=60)
resp.raise_for_status()
output_path = os.path.join(os.path.dirname(__file__), "output", "test_output.mp4")
with open(output_path, "wb") as f:
    f.write(resp.content)
size_kb = os.path.getsize(output_path) / 1024
print(f"    Saved: {output_path} ({size_kb:.1f} KB)")

# Save metadata
meta = {"video_id": video_id, "url": video_url, "resolution": result.get("size"), "model": result.get("model")}
with open(os.path.join(os.path.dirname(__file__), "output", "meta.json"), "w", encoding="utf-8") as f:
    json.dump(meta, f, indent=2, ensure_ascii=False)

print("\n=== TEST PASSED ===")