"""Tests for the PyAV-backed video compatibility surface.""" from __future__ import annotations import os import subprocess import sys from pathlib import Path import av import numpy as np import pytest from supervision import _cv2 from supervision._cv2._video import ( _mux_audio, _VideoCapture, _VideoWriter, ) def _write_video(path: Path, values: list[int], fps: int = 5) -> None: """Write a small deterministic MPEG-4 video for fallback tests.""" container = av.open(str(path), mode="w") stream = container.add_stream("mpeg4", rate=fps) stream.width = 16 stream.height = 16 stream.pix_fmt = "yuv420p" try: for value in values: frame = av.VideoFrame.from_ndarray( np.full((16, 16, 3), value, dtype=np.uint8), format="bgr24" ) for packet in stream.encode(frame): container.mux(packet) for packet in stream.encode(): container.mux(packet) finally: container.close() def _write_video_with_audio(path: Path, frame_count: int = 5, fps: int = 5) -> None: """Write a short video with one AAC audio stream for remux tests.""" container = av.open(str(path), mode="w") video_stream = container.add_stream("mpeg4", rate=fps) video_stream.width = 16 video_stream.height = 16 video_stream.pix_fmt = "yuv420p" audio_stream = container.add_stream("aac", rate=8_000) audio_stream.layout = "mono" try: for value in range(frame_count): frame = av.VideoFrame.from_ndarray( np.full((16, 16, 3), value * 20, dtype=np.uint8), format="bgr24" ) for packet in video_stream.encode(frame): container.mux(packet) samples = np.zeros((1, 8_000), dtype=np.int16) audio_frame = av.AudioFrame.from_ndarray(samples, format="s16", layout="mono") audio_frame.sample_rate = 8_000 audio_frame.pts = 0 for packet in audio_stream.encode(audio_frame): container.mux(packet) for packet in video_stream.encode(): container.mux(packet) for packet in audio_stream.encode(): container.mux(packet) finally: container.close() def _run_without_opencv(source: str) -> None: """Run a Python snippet with cv2 imports blocked.""" env = os.environ.copy() source_path = str(Path(__file__).resolve().parents[2] / "src") env["PYTHONPATH"] = os.pathsep.join( filter(None, (source_path, env.get("PYTHONPATH"))) ) result = subprocess.run( [sys.executable, "-c", source], check=False, capture_output=True, text=True, env=env, timeout=60, ) assert result.returncode == 0, result.stderr def test_video_module_uses_required_pyav_dependency() -> None: """The video fallback imports required PyAV directly without a loader.""" from supervision._cv2 import _video assert _video.av.__name__ == "av" assert not hasattr(_video, "_load_av") def test_fallback_capture_reports_metadata_and_supports_exact_seek( tmp_path: Path, ) -> None: """Fallback capture exposes metadata and starts decoding at requested frames.""" source_path = tmp_path / "source.mp4" _write_video(source_path, [0, 40, 80, 120, 160]) capture = _VideoCapture(str(source_path)) assert capture.isOpened() assert capture.get(_cv2.CAP_PROP_FRAME_WIDTH) == 16 assert capture.get(_cv2.CAP_PROP_FRAME_HEIGHT) == 16 assert capture.get(_cv2.CAP_PROP_FPS) == pytest.approx(5.0) assert capture.get(_cv2.CAP_PROP_FRAME_COUNT) == 5 assert capture.set(_cv2.CAP_PROP_POS_FRAMES, 3) success, frame = capture.read() capture.release() assert success assert frame is not None assert frame.shape == (16, 16, 3) assert float(frame.mean()) == pytest.approx(120.0, abs=20.0) assert not capture.isOpened() def test_fallback_writer_default_codec_round_trips(tmp_path: Path) -> None: """The guaranteed mp4v fallback writer creates a readable video.""" target_path = tmp_path / "target.mp4" fourcc = _cv2.VideoWriter_fourcc(*"mp4v") writer = _VideoWriter(str(target_path), fourcc, 5.0, (16, 16)) assert writer.isOpened() for value in [0, 40, 80]: writer.write(np.full((16, 16, 3), value, dtype=np.uint8)) writer.release() capture = _VideoCapture(str(target_path)) frames = [] while True: success, frame = capture.read() if not success: break assert frame is not None frames.append(frame) capture.release() assert len(frames) == 3 assert target_path.stat().st_size > 0 def test_fallback_video_works_when_opencv_is_blocked(tmp_path: Path) -> None: """Production video APIs use PyAV when cv2 cannot be imported.""" source_path = tmp_path / "source.mp4" target_path = tmp_path / "target.mp4" _write_video(source_path, [0, 40, 80]) source = f""" import sys class BlockCv2: def find_spec(self, fullname, path=None, target=None): if fullname == "cv2": raise ModuleNotFoundError("blocked for test") return None sys.meta_path.insert(0, BlockCv2()) from supervision import _cv2 from supervision.utils.video import VideoInfo, VideoSink, get_video_frames_generator assert _cv2._IS_CV2_AVAILABLE is False info = VideoInfo.from_video_path({str(source_path)!r}) assert (info.width, info.height, info.total_frames) == (16, 16, 3) frames = list(get_video_frames_generator({str(source_path)!r}, start=1, end=3)) assert len(frames) == 2 with VideoSink({str(target_path)!r}, info) as sink: for frame in frames: sink.write_frame(frame) assert _cv2.VideoCapture({str(target_path)!r}).isOpened() """ _run_without_opencv(source) def test_mux_audio_remuxes_first_audio_stream_and_truncates_to_video( tmp_path: Path, ) -> None: """Audio remuxing uses PyAV and keeps the processed video duration.""" source_path = tmp_path / "source_with_audio.mp4" target_path = tmp_path / "target.mp4" _write_video_with_audio(source_path, frame_count=5) _write_video(target_path, [0, 40, 80], fps=5) _mux_audio(str(source_path), str(target_path)) output = av.open(str(target_path)) try: assert len(output.streams.video) == 1 assert len(output.streams.audio) == 1 assert output.streams.video[0].frames == 3 finally: output.close() def test_mux_audio_leaves_target_unchanged_on_failure(tmp_path: Path) -> None: """A failed PyAV remux never replaces the existing target file.""" target_path = tmp_path / "target.mp4" _write_video(target_path, [0, 40, 80]) original = target_path.read_bytes() _mux_audio(str(tmp_path / "missing.mp4"), str(target_path)) assert target_path.read_bytes() == original