code-test-drone/utils/tests/ground_station/video_source_test.py

"""video_source_test.py —— RTSP→MJPEG 视频源单元测试（注入假采集器，不依赖真流）。

覆盖：拉帧→JPEG 编码缓存、宽度限制（等比缩放）、断线重连（重新调用工厂）、
打开失败退避重试、online 判定、MJPEG 分块格式。
"""

import sys
import time
from pathlib import Path

import numpy as np

_ROOT = Path(__file__).resolve().parents[3]   # utils/tests/ground_station → 仓库根
if str(_ROOT) not in sys.path:
    sys.path.insert(0, str(_ROOT))

from ground_station.video_source import RtspVideoSource  # noqa: E402


class FakeCapture:
    """假采集器：可注入读失败次数；帧为纯色小图。"""

    def __init__(self, fail_reads=0, height=60, width=120):
        self.fail_reads = fail_reads
        self._frame = np.full((height, width, 3), 128, dtype=np.uint8)
        self.released = False

    def isOpened(self):
        return True

    def read(self):
        if self.fail_reads > 0:
            self.fail_reads -= 1
            return False, None
        return True, self._frame

    def release(self):
        self.released = True


def _wait(pred, timeout_s=4.0, step=0.02):
    t0 = time.monotonic()
    while time.monotonic() - t0 < timeout_s:
        if pred():
            return True
        time.sleep(step)
    return False


def _decode(jpeg):
    import cv2
    return cv2.imdecode(np.frombuffer(jpeg, dtype=np.uint8), cv2.IMREAD_COLOR)


def test_frames_encoded_and_online():
    src = RtspVideoSource("rtsp://fake", width=120, fps=50,
                          capture_factory=lambda: FakeCapture())
    assert src.online is False                  # 尚无帧 → 离线
    src.start()
    try:
        assert _wait(lambda: src.latest()[1] is not None)
        seq, jpeg = src.latest()
        assert seq >= 1
        assert jpeg[:2] == b"\xff\xd8"          # JPEG SOI 头
        assert _decode(jpeg).shape[1] == 120    # 宽度不超限
        assert src.online is True
    finally:
        src.stop()


def test_downscale_to_max_width():
    src = RtspVideoSource("rtsp://fake", width=60, fps=50,
                          capture_factory=lambda: FakeCapture(height=60, width=120))
    src.start()
    try:
        assert _wait(lambda: src.latest()[1] is not None)
        img = _decode(src.latest()[1])
        assert img.shape[1] == 60 and img.shape[0] == 30   # 等比缩到 60 宽
    finally:
        src.stop()


def test_reconnect_after_read_failure():
    opened = []

    def factory():
        cap = FakeCapture(fail_reads=(10 ** 6 if not opened else 0))
        opened.append(cap)
        return cap

    src = RtspVideoSource("rtsp://fake", width=120, fps=50,
                          capture_factory=factory)
    src.start()
    try:
        assert _wait(lambda: src.latest()[1] is not None, timeout_s=5.0)
        assert len(opened) >= 2                 # 首采集器读失败 → 重新打开
        assert opened[0].released is True
    finally:
        src.stop()


def test_open_failure_keeps_retrying():
    calls = []

    class ClosedCapture:
        def isOpened(self):
            return False

        def release(self):
            pass

    def factory():
        calls.append(1)
        return ClosedCapture()

    src = RtspVideoSource("rtsp://fake", capture_factory=factory)
    src.start()
    try:
        assert _wait(lambda: len(calls) >= 2, timeout_s=4.0)  # 1s 退避后重试
        assert src.latest()[1] is None
    finally:
        src.stop()


def test_mjpeg_stream_chunks():
    src = RtspVideoSource("rtsp://fake", width=120, fps=50,
                          capture_factory=lambda: FakeCapture())
    src.start()
    try:
        assert _wait(lambda: src.latest()[1] is not None)
        gen = src.mjpeg_stream()
        chunk = next(gen)
        assert chunk.startswith(b"--frame\r\nContent-Type: image/jpeg\r\n")
        assert b"Content-Length: " in chunk
        assert chunk.endswith(b"\r\n")
        gen.close()
    finally:
        src.stop()