tests/test_frames.py

"""Tests for fpgm.pipeline.frames.ClipFrameSource, including its ``stride`` option.

A tiny synthetic video is written with cv2 into ``tmp_path`` -- no real DROID mp4
needed. Each frame is a solid colour encoding its own index, so extraction can be
checked by decoding the colour back rather than trusting byte-identity.
"""

from __future__ import annotations

import cv2
import numpy as np
import pytest

from fpgm.pipeline.frames import ClipFrameSource


def _write_synthetic_video(path, n_frames: int, size: tuple[int, int] = (32, 24)) -> None:
    """Write ``n_frames`` solid-colour BGR frames; frame ``i`` is colour ``(10*i, ...)``.

    MJPG (intra-frame only, unlike mp4v's inter-frame prediction) is used
    deliberately: a temporally-predicted codec smears a near-constant colour ramp
    across frames, which would make this test measure JPEG/MPEG artifacts instead
    of ``ClipFrameSource``'s own frame selection.
    """
    width, height = size
    writer = cv2.VideoWriter(
        str(path), cv2.VideoWriter_fourcc(*"MJPG"), 10.0, (width, height)
    )
    try:
        for i in range(n_frames):
            frame = np.full((height, width, 3), fill_value=min(10 * i, 255), dtype=np.uint8)
            writer.write(frame)
    finally:
        writer.release()


def _frame_value(bgr: np.ndarray) -> int:
    return round(int(bgr[0, 0, 0]) / 10.0)


class TestClipFrameSourceStride:
    def test_default_stride_extracts_consecutive_frames(self, tmp_path):
        video_path = tmp_path / "video.avi"
        _write_synthetic_video(video_path, n_frames=20)

        with ClipFrameSource(video_path, start_frame=2, end_frame=8) as frames:
            assert frames.n_frames == 6
            values = [_frame_value(frames.read(i)) for i in range(frames.n_frames)]
        assert values == [2, 3, 4, 5, 6, 7]

    def test_stride_two_extracts_every_other_frame(self, tmp_path):
        video_path = tmp_path / "video.avi"
        _write_synthetic_video(video_path, n_frames=30)

        # Mirrors the real bug: annotation clip "5:11" (6 frames) maps to video
        # frames [10, 22) with stride 2.
        with ClipFrameSource(video_path, start_frame=10, end_frame=22, stride=2) as frames:
            assert frames.n_frames == 6
            values = [_frame_value(frames.read(i)) for i in range(frames.n_frames)]
        assert values == [10, 12, 14, 16, 18, 20]

    def test_stride_matches_python_range_semantics(self, tmp_path):
        video_path = tmp_path / "video.avi"
        _write_synthetic_video(video_path, n_frames=25)

        start, end, stride = 3, 19, 4
        expected = list(range(start, end, stride))
        with ClipFrameSource(video_path, start, end, stride=stride) as frames:
            assert frames.n_frames == len(expected)
            values = [_frame_value(frames.read(i)) for i in range(frames.n_frames)]
        assert values == expected

    def test_stride_truncates_gracefully_past_video_end(self, tmp_path):
        video_path = tmp_path / "video.avi"
        _write_synthetic_video(video_path, n_frames=15)

        # Requested end (30) is past the 15-frame video; extraction should still
        # succeed with whatever strided frames actually exist, not raise.
        with ClipFrameSource(video_path, start_frame=10, end_frame=30, stride=2) as frames:
            values = [_frame_value(frames.read(i)) for i in range(frames.n_frames)]
        assert values == [10, 12, 14]

    def test_invalid_stride_rejected(self, tmp_path):
        video_path = tmp_path / "video.avi"
        _write_synthetic_video(video_path, n_frames=5)
        with pytest.raises(ValueError):
            ClipFrameSource(video_path, 0, 5, stride=0)