silver_pose/v1/tests/test_video_source.py

import cv2
import numpy as np
import pytest

from v1.video_source import SourceMode, SourceStatus, VideoSource


def _write_sample_video(path):
    writer = cv2.VideoWriter(
        str(path), cv2.VideoWriter_fourcc(*"MJPG"), 10.0, (32, 24)
    )
    assert writer.isOpened()
    for value in (20, 80, 140):
        writer.write(np.full((24, 32, 3), value, dtype=np.uint8))
    writer.release()


def test_file_source_emits_monotonic_timestamps(tmp_path):
    sample_video = tmp_path / "sample.avi"
    _write_sample_video(sample_video)
    source = VideoSource(sample_video, mode=SourceMode.REPLAY)

    first = source.read(now=10.0)
    second = source.read(now=10.1)

    assert first.status is SourceStatus.CONNECTED
    assert second.status is SourceStatus.CONNECTED
    assert first.image is not None
    assert second.image is not None
    assert first.timestamp_monotonic < second.timestamp_monotonic


def test_missing_source_returns_error_state_without_frame(tmp_path):
    source = VideoSource(tmp_path / "missing.avi", mode=SourceMode.REPLAY)

    packet = source.read(now=1.0)

    assert packet.status is SourceStatus.ERROR
    assert packet.image is None
    assert packet.error


class _ClosedCapture:
    def isOpened(self):
        return False

    def read(self):
        return False, None

    def get(self, _property_id):
        return 0.0

    def release(self):
        pass


class _OpenCapture:
    def __init__(self):
        self._read_count = 0

    def isOpened(self):
        return True

    def read(self):
        self._read_count += 1
        return True, np.zeros((8, 8, 3), dtype=np.uint8)

    def get(self, _property_id):
        return self._read_count * 100.0

    def release(self):
        pass


def test_reconnect_waits_then_reopens_with_bounded_backoff():
    captures = [_ClosedCapture(), _OpenCapture()]

    source = VideoSource(
        "demo-source",
        mode=SourceMode.STREAM,
        retry_initial_seconds=2.0,
        retry_max_seconds=2.0,
        capture_factory=lambda _source: captures.pop(0),
    )

    first = source.read(now=5.0)
    waiting = source.read(now=6.0)
    recovered = source.read(now=7.0)

    assert first.status is SourceStatus.RETRYING
    assert waiting.status is SourceStatus.RETRYING
    assert recovered.status is SourceStatus.CONNECTED
    assert recovered.image is not None


class _NegativeFirstTimestampCapture:
    def __init__(self):
        self._read_count = 0

    def isOpened(self):
        return True

    def read(self):
        self._read_count += 1
        return True, np.zeros((8, 8, 3), dtype=np.uint8)

    def get(self, property_id):
        if property_id == cv2.CAP_PROP_POS_MSEC:
            return -33.0 if self._read_count == 1 else 33.333333333333336
        if property_id == cv2.CAP_PROP_POS_FRAMES:
            return float(self._read_count)
        if property_id == cv2.CAP_PROP_FPS:
            return 30.0
        return 0.0

    def release(self):
        pass


def test_negative_first_timestamp_falls_back_to_frame_index_and_fps():
    source = VideoSource(
        "demo-source",
        mode=SourceMode.REPLAY,
        capture_factory=lambda _source: _NegativeFirstTimestampCapture(),
    )

    first = source.read(now=10.0)
    second = source.read(now=10.1)

    assert first.timestamp_monotonic == 0.0
    assert second.timestamp_monotonic == pytest.approx(1.0 / 30.0)


def test_replay_source_reports_eof_without_restarting(tmp_path):
    sample_video = tmp_path / "sample.avi"
    _write_sample_video(sample_video)
    source = VideoSource(sample_video, mode=SourceMode.REPLAY)

    for timestamp in (0.0, 0.1, 0.2):
        assert source.read(now=timestamp).status is SourceStatus.CONNECTED
    eof = source.read(now=0.3)
    still_eof = source.read(now=0.4)

    assert eof.status is SourceStatus.EOF
    assert still_eof.status is SourceStatus.EOF


def test_stream_source_uses_read_clock_instead_of_capture_timestamp():
    source = VideoSource(
        "rtsp://not-a-real-address",
        mode=SourceMode.STREAM,
        capture_factory=lambda _source: _OpenCapture(),
    )

    packet = source.read(now=42.0)

    assert packet.status is SourceStatus.CONNECTED
    assert packet.timestamp_monotonic == 42.0