video-ai-analysis/tests/test_clips.py

import json
import tempfile
import unittest
from pathlib import Path

from video_ai_analysis_poc.clips import build_clip_records, build_clip_records_from_manifest


class ClipTests(unittest.TestCase):
    def test_build_clip_records_uniformly_samples_frames_per_clip(self):
        frames = [
            {
                "video_id": "video-abc",
                "frame_id": f"video-abc_f{index + 1:06d}",
                "frame_path": f"frames/video-abc/{index + 1:06d}.jpg",
                "offset_seconds": float(index),
                "timecode": f"00:00:{index:02d}",
                "pts_time": float(index),
                "status": "sampled",
            }
            for index in range(10)
        ]

        clips = build_clip_records(
            frames,
            {
                "length_seconds": 10,
                "stride_seconds": 10,
                "frames_per_clip": 4,
                "min_frames_per_clip": 2,
            },
        )

        self.assertEqual(len(clips), 1)
        self.assertEqual(clips[0]["clip_id"], "video-abc_c000001")
        self.assertEqual(clips[0]["clip_start_seconds"], 0.0)
        self.assertEqual(clips[0]["clip_end_seconds"], 10.0)
        self.assertEqual(
            [frame["offset_seconds"] for frame in clips[0]["frame_times"]],
            [0.0, 3.0, 6.0, 9.0],
        )
        self.assertEqual(clips[0]["status"], "pending")
        self.assertEqual(clips[0]["retry_count"], 0)
        self.assertIsNone(clips[0]["last_error"])

    def test_tail_clip_end_is_truncated_to_last_frame_interval(self):
        frames = [
            {
                "video_id": "video-abc",
                "frame_id": f"video-abc_f{index + 1:06d}",
                "frame_path": f"frames/video-abc/{index + 1:06d}.jpg",
                "offset_seconds": float(index),
                "timecode": f"00:00:{index:02d}",
                "pts_time": float(index),
                "status": "sampled",
            }
            for index in range(15)
        ]

        clips = build_clip_records(
            frames,
            {
                "length_seconds": 10,
                "stride_seconds": 10,
                "frames_per_clip": 8,
                "min_frames_per_clip": 4,
            },
        )

        self.assertEqual(len(clips), 2)
        self.assertEqual(clips[1]["clip_start_seconds"], 10.0)
        self.assertEqual(clips[1]["clip_end_seconds"], 15.0)
        self.assertEqual(clips[1]["clip_end_timecode"], "00:00:15")

    def test_build_clip_records_adds_beijing_time_range_and_frame_times(self):
        frames = [
            {
                "video_id": "video-abc",
                "frame_id": f"video-abc_f{index + 1:06d}",
                "frame_path": f"frames/video-abc/{index + 1:06d}.jpg",
                "offset_seconds": float(index),
                "timecode": f"00:00:{index:02d}",
                "pts_time": float(index),
                "beijing_time": f"2026-06-15 07:00:{index:02d}",
                "status": "sampled",
            }
            for index in range(10)
        ]

        clips = build_clip_records(
            frames,
            {
                "length_seconds": 10,
                "stride_seconds": 10,
                "frames_per_clip": 4,
                "min_frames_per_clip": 2,
            },
        )

        self.assertEqual(clips[0]["clip_start_beijing_time"], "2026-06-15 07:00:00")
        self.assertEqual(clips[0]["clip_end_beijing_time"], "2026-06-15 07:00:10")
        self.assertEqual(
            [frame["beijing_time"] for frame in clips[0]["frame_times"]],
            [
                "2026-06-15 07:00:00",
                "2026-06-15 07:00:03",
                "2026-06-15 07:00:06",
                "2026-06-15 07:00:09",
            ],
        )

    def test_build_clip_records_from_manifest_skips_failed_frames_and_writes_jsonl(self):
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            frame_manifest = root / "frame_manifest.jsonl"
            clip_manifest = root / "clip_manifest.jsonl"
            records = [
                {
                    "video_id": "video-abc",
                    "frame_id": f"video-abc_f{index + 1:06d}",
                    "frame_path": f"frames/video-abc/{index + 1:06d}.jpg",
                    "offset_seconds": float(index),
                    "timecode": f"00:00:{index:02d}",
                    "pts_time": float(index),
                    "status": "sampled",
                }
                for index in range(4)
            ]
            records.append(
                {
                    "video_id": "video-abc",
                    "frame_id": None,
                    "frame_path": None,
                    "offset_seconds": None,
                    "timecode": None,
                    "pts_time": None,
                    "status": "sample_failed",
                    "last_error": "bad decode",
                }
            )
            frame_manifest.write_text(
                "\n".join(json.dumps(record, sort_keys=True) for record in records) + "\n",
                encoding="utf-8",
            )

            clips = build_clip_records_from_manifest(
                frame_manifest,
                clip_manifest,
                {
                    "length_seconds": 10,
                    "stride_seconds": 10,
                    "frames_per_clip": 8,
                    "min_frames_per_clip": 4,
                },
            )

            self.assertEqual(len(clips), 1)
            self.assertEqual(len(clips[0]["frame_times"]), 4)
            persisted = [
                json.loads(line)
                for line in clip_manifest.read_text(encoding="utf-8").splitlines()
            ]
            self.assertEqual(persisted, clips)


if __name__ == "__main__":
    unittest.main()