返回 ViMax
test_openrouter_video_generator.py
根目录 / tests / test_openrouter_video_generator.py
1 import os
2 import tempfile
3 import unittest
4 from pathlib import Path
5 from unittest.mock import patch
6
7 from PIL import Image
8
9 from tools.video_generator_openrouter_api import VideoGeneratorOpenRouterAPI
10 from interfaces.video_output import VideoOutput
11
12
13 class OpenRouterVideoGeneratorTests(unittest.IsolatedAsyncioTestCase):
14 async def test_default_duration_is_eight_seconds(self):
15 captured = {}
16
17 async def fake_post_json(url, *, headers, payload, timeout, hard_timeout_seconds):
18 captured["payload"] = payload
19 return 200, {"id": "job-1", "polling_url": "/videos/job-1", "status": "queued"}
20
21 async def fake_get_json(url, *, headers, timeout, hard_timeout_seconds):
22 return 200, {"status": "completed", "unsigned_urls": ["https://cdn.example/out.mp4"]}
23
24 async def fake_get_bytes(url, *, headers, timeout, hard_timeout_seconds):
25 return 200, b"video"
26
27 async def fake_sleep(seconds):
28 return None
29
30 generator = VideoGeneratorOpenRouterAPI(api_key="test-key", model="google/veo-3.1-lite")
31 with patch.dict(os.environ, {}, clear=True), \
32 patch("tools.video_generator_openrouter_api._post_json", fake_post_json), \
33 patch("tools.video_generator_openrouter_api._get_json", fake_get_json), \
34 patch("tools.video_generator_openrouter_api._get_bytes", fake_get_bytes), \
35 patch("tools.video_generator_openrouter_api.asyncio.sleep", fake_sleep):
36 output = await generator.generate_single_video(prompt="hello")
37
38 self.assertIsInstance(output, VideoOutput)
39 self.assertEqual(captured["payload"]["duration"], 8)
40 self.assertEqual(captured["payload"]["model"], "google/veo-3.1-lite")
41
42 async def test_seedance_fast_uses_supported_openrouter_payload(self):
43 captured = {}
44
45 async def fake_post_json(url, *, headers, payload, timeout, hard_timeout_seconds):
46 captured["payload"] = payload
47 return 200, {"id": "job-2", "polling_url": "/videos/job-2", "status": "queued"}
48
49 async def fake_get_json(url, *, headers, timeout, hard_timeout_seconds):
50 return 200, {"status": "completed", "unsigned_urls": ["https://cdn.example/seedance.mp4"]}
51
52 async def fake_get_bytes(url, *, headers, timeout, hard_timeout_seconds):
53 return 200, b"seedance-video"
54
55 async def fake_sleep(seconds):
56 return None
57
58 with tempfile.TemporaryDirectory() as tmp:
59 first_frame = Path(tmp) / "first.png"
60 last_frame = Path(tmp) / "last.png"
61 Image.new("RGB", (16, 9), "blue").save(first_frame)
62 Image.new("RGB", (16, 9), "red").save(last_frame)
63 generator = VideoGeneratorOpenRouterAPI(api_key="test-key", model="bytedance/seedance-2.0-fast")
64 with patch.dict(os.environ, {}, clear=True), \
65 patch("tools.video_generator_openrouter_api._post_json", fake_post_json), \
66 patch("tools.video_generator_openrouter_api._get_json", fake_get_json), \
67 patch("tools.video_generator_openrouter_api._get_bytes", fake_get_bytes), \
68 patch("tools.video_generator_openrouter_api.asyncio.sleep", fake_sleep):
69 output = await generator.generate_single_video(
70 prompt="a cinematic walk",
71 reference_image_paths=[str(first_frame), str(last_frame)],
72 )
73
74 payload = captured["payload"]
75 self.assertIsInstance(output, VideoOutput)
76 self.assertEqual(payload["model"], "bytedance/seedance-2.0-fast")
77 self.assertEqual(payload["duration"], 8)
78 self.assertEqual(payload["resolution"], "720p")
79 self.assertEqual(payload["aspect_ratio"], "16:9")
80 self.assertTrue(payload["generate_audio"])
81 self.assertEqual(
82 [frame["frame_type"] for frame in payload["frame_images"]],
83 ["first_frame", "last_frame"],
84 )
85
86
87 if __name__ == "__main__":
88 unittest.main()
89
89 lines PYTHON