| 1 | import tempfile |
| 2 | import unittest |
| 3 | from pathlib import Path |
| 4 | |
| 5 | from interfaces import Camera, ShotBriefDescription, ShotDescription |
| 6 | from pipelines.script2video_pipeline import Script2VideoPipeline, _group_shots_into_cameras |
| 7 | |
| 8 | |
| 9 | class FlakyCameraImageGenerator: |
| 10 | def __init__(self): |
| 11 | self.calls = 0 |
| 12 | |
| 13 | async def construct_camera_tree(self, cameras, shot_descs): |
| 14 | self.calls += 1 |
| 15 | if self.calls == 1: |
| 16 | return ["not-a-camera"] |
| 17 | return cameras |
| 18 | |
| 19 | |
| 20 | class Script2VideoPipelineGuardTests(unittest.IsolatedAsyncioTestCase): |
| 21 | def test_group_shots_into_cameras_does_not_use_camera_idx_as_list_index(self): |
| 22 | shots = [ |
| 23 | ShotDescription(idx=0, is_last=False, cam_idx=2, visual_desc="a", variation_type="small", variation_reason="same", ff_desc="a", ff_vis_char_idxs=[], lf_desc="a", lf_vis_char_idxs=[], motion_desc="a", audio_desc="none"), |
| 24 | ShotDescription(idx=1, is_last=True, cam_idx=5, visual_desc="b", variation_type="small", variation_reason="same", ff_desc="b", ff_vis_char_idxs=[], lf_desc="b", lf_vis_char_idxs=[], motion_desc="b", audio_desc="none"), |
| 25 | ShotDescription(idx=2, is_last=True, cam_idx=2, visual_desc="c", variation_type="small", variation_reason="same", ff_desc="c", ff_vis_char_idxs=[], lf_desc="c", lf_vis_char_idxs=[], motion_desc="c", audio_desc="none"), |
| 26 | ] |
| 27 | cameras = _group_shots_into_cameras(shots) |
| 28 | self.assertEqual([camera.idx for camera in cameras], [2, 5]) |
| 29 | self.assertEqual(cameras[0].active_shot_idxs, [0, 2]) |
| 30 | self.assertEqual(cameras[1].active_shot_idxs, [1]) |
| 31 | |
| 32 | async def test_plan_text_artifacts_retries_bad_camera_tree_schema(self): |
| 33 | with tempfile.TemporaryDirectory() as tmp: |
| 34 | pipeline = Script2VideoPipeline(chat_model=object(), image_generator=object(), video_generator=object(), working_dir=tmp) |
| 35 | pipeline.camera_image_generator = FlakyCameraImageGenerator() |
| 36 | |
| 37 | async def design_storyboard(script, characters, user_requirement, quiet=False): |
| 38 | return [{"idx": 0, "is_last": True, "cam_idx": 3, "visual_desc": "wide shot", "audio_desc": "waves"}] |
| 39 | |
| 40 | async def decompose_visual_descriptions(shot_brief_descriptions, characters, quiet=False): |
| 41 | return [{"idx": 0, "is_last": True, "cam_idx": 3, "visual_desc": "wide shot", "variation_type": "small", "variation_reason": "simple", "ff_desc": "start", "ff_vis_char_idxs": [], "lf_desc": "end", "lf_vis_char_idxs": [], "motion_desc": "walk", "audio_desc": "waves"}] |
| 42 | |
| 43 | pipeline.design_storyboard = design_storyboard |
| 44 | pipeline.decompose_visual_descriptions = decompose_visual_descriptions |
| 45 | events = [] |
| 46 | result = await pipeline.plan_text_artifacts( |
| 47 | "script", |
| 48 | "req", |
| 49 | "style", |
| 50 | characters=[{"idx": 0, "identifier_in_scene": "Man", "is_visible": True, "static_features": "adult", "dynamic_features": "coat"}], |
| 51 | progress=lambda stage, message, metadata=None: events.append(stage), |
| 52 | quiet=True, |
| 53 | ) |
| 54 | |
| 55 | self.assertEqual(pipeline.camera_image_generator.calls, 2) |
| 56 | self.assertIn("construct_camera_tree_retry", events) |
| 57 | self.assertEqual(result["camera_tree"][0].idx, 3) |
| 58 | self.assertTrue((Path(tmp) / "camera_tree.json").exists()) |
| 59 | |
| 60 | |
| 61 | if __name__ == "__main__": |
| 62 | unittest.main() |
| 63 |