silver_pose/v1/tests/test_video_source.py

import time

import cv2
import numpy as np
import pytest

from v1.video_source import SourceMode, SourceStatus, VideoSource, _FrameGrabber


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 _FlakyCapture:
    def __init__(self, read_results):
        self._read_results = list(read_results)

    def isOpened(self):
        return True

    def read(self):
        ok = self._read_results.pop(0) if self._read_results else False
        return (True, np.zeros((8, 8, 3), dtype=np.uint8)) if ok else (False, None)

    def get(self, _property_id):
        return 0.0

    def release(self):
        pass


def test_stream_recovers_after_mid_stream_drop():
    captures = [_FlakyCapture([True]), _OpenCapture()]

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

    connected = source.read(now=0.0)
    dropped = source.read(now=1.0)
    waiting = source.read(now=1.5)
    recovered = source.read(now=3.0)

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


class _ScriptedCapture:
    def __init__(self, frames, fail_after=False):
        self._frames = list(frames)
        self._fail_after = fail_after

    def isOpened(self):
        return True

    def read(self):
        if self._frames:
            return True, self._frames.pop(0)
        return False, None

    def get(self, _property_id):
        return 0.0

    def release(self):
        pass


def test_frame_grabber_keeps_only_the_latest_frame():
    frame_a = np.full((4, 4, 3), 1, dtype=np.uint8)
    frame_b = np.full((4, 4, 3), 2, dtype=np.uint8)
    frame_c = np.full((4, 4, 3), 3, dtype=np.uint8)
    grabber = _FrameGrabber(_ScriptedCapture([frame_a, frame_b, frame_c]))

    assert grabber._pump() is True
    assert grabber._pump() is True
    assert grabber._pump() is True
    image, failed = grabber.take_latest()

    assert failed is False
    assert image is frame_c  # older frames dropped


def test_frame_grabber_marks_failed_on_read_failure():
    grabber = _FrameGrabber(_ScriptedCapture([np.zeros((4, 4, 3), dtype=np.uint8)]))

    assert grabber._pump() is True
    assert grabber._pump() is False
    _, failed = grabber.take_latest()

    assert failed is True


def test_stream_with_drop_stale_returns_a_connected_frame():
    source = VideoSource(
        "rtsp://demo",
        mode=SourceMode.STREAM,
        drop_stale=True,
        capture_factory=lambda _s: _OpenCapture(),
    )
    try:
        packet = None
        for _ in range(40):
            packet = source.read()
            if packet.status is SourceStatus.CONNECTED and packet.image is not None:
                break
            time.sleep(0.02)
        assert packet is not None
        assert packet.status is SourceStatus.CONNECTED
        assert packet.image is not None
    finally:
        source.close()


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