"""Round-trip tests for VideoSink and ImageSink.""" from pathlib import Path import numpy as np import pytest import supervision as sv import supervision.utils.video as video_utils from supervision.utils.video import VideoInfo class TestImageSink: """ImageSink saves images to disk using the configured name pattern.""" def test_saves_single_image(self, tmp_path: Path) -> None: """A single saved frame produces exactly one file.""" image = np.zeros((64, 64, 3), dtype=np.uint8) with sv.ImageSink(target_dir_path=str(tmp_path), overwrite=True) as sink: sink.save_image(image=image) files = list(tmp_path.iterdir()) assert len(files) == 1 def test_saves_multiple_images(self, tmp_path: Path) -> None: """N consecutive save_image calls create N distinct files.""" image = np.zeros((64, 64, 3), dtype=np.uint8) n = 5 with sv.ImageSink(target_dir_path=str(tmp_path), overwrite=True) as sink: for _ in range(n): sink.save_image(image=image) files = sorted(tmp_path.iterdir()) assert len(files) == n def test_custom_name_pattern(self, tmp_path: Path) -> None: """Custom image_name_pattern is applied to each saved file.""" image = np.zeros((8, 8, 3), dtype=np.uint8) with sv.ImageSink( target_dir_path=str(tmp_path), overwrite=True, image_name_pattern="frame_{:03d}.jpg", ) as sink: sink.save_image(image=image) names = [f.name for f in tmp_path.iterdir()] assert names == ["frame_000.jpg"] def test_overwrite_false_reuses_existing_dir(self, tmp_path: Path) -> None: """overwrite=False keeps existing directory contents intact.""" existing = tmp_path / "existing" existing.mkdir() sentinel = existing / "keep.txt" sentinel.write_text("keep") image = np.zeros((8, 8, 3), dtype=np.uint8) with sv.ImageSink(target_dir_path=str(existing), overwrite=False) as sink: sink.save_image(image=image) assert sentinel.exists(), "pre-existing file should not be deleted" assert len(list(existing.iterdir())) == 2 def test_overwrite_true_clears_existing_dir(self, tmp_path: Path) -> None: """overwrite=True removes pre-existing files before writing.""" target = tmp_path / "target" target.mkdir() sentinel = target / "old.txt" sentinel.write_text("old") image = np.zeros((8, 8, 3), dtype=np.uint8) with sv.ImageSink(target_dir_path=str(target), overwrite=True) as sink: sink.save_image(image=image) assert not sentinel.exists(), "overwrite=True must clear pre-existing files" class TestVideoSink: """VideoSink writes valid video frames to a file.""" def test_write_frame_outside_context_raises(self, tmp_path: Path) -> None: """write_frame requires an open VideoSink context.""" target = str(tmp_path / "out.avi") info = VideoInfo(width=64, height=64, fps=1, total_frames=None) sink = sv.VideoSink(target_path=target, video_info=info, codec="MJPG") with pytest.raises(RuntimeError, match="open VideoSink context"): sink.write_frame(frame=np.zeros((64, 64, 3), dtype=np.uint8)) def test_enter_raises_when_writer_does_not_open( self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path ) -> None: """VideoSink surfaces VideoWriter open failures immediately.""" class ClosedWriter: def isOpened(self) -> bool: return False def release(self) -> None: pass monkeypatch.setattr( video_utils.cv2, "VideoWriter", lambda *args: ClosedWriter() ) target = str(tmp_path / "out.avi") info = VideoInfo(width=64, height=64, fps=1, total_frames=None) with pytest.raises(RuntimeError, match="Could not open video writer"): with sv.VideoSink(target_path=target, video_info=info, codec="MJPG"): pass def test_creates_output_file(self, tmp_path: Path) -> None: """VideoSink creates a non-empty file at target_path.""" target = str(tmp_path / "out.avi") info = VideoInfo(width=64, height=64, fps=1, total_frames=None) frame = np.zeros((64, 64, 3), dtype=np.uint8) with sv.VideoSink(target_path=target, video_info=info, codec="MJPG") as sink: sink.write_frame(frame=frame) assert Path(target).exists() assert Path(target).stat().st_size > 0 def test_writes_multiple_frames(self, tmp_path: Path) -> None: """Writing N frames produces a strictly larger file than writing one frame.""" info = VideoInfo(width=64, height=64, fps=5, total_frames=None) target_one = str(tmp_path / "one.avi") with sv.VideoSink( target_path=target_one, video_info=info, codec="MJPG" ) as sink: sink.write_frame(frame=np.zeros((64, 64, 3), dtype=np.uint8)) target_many = str(tmp_path / "many.avi") with sv.VideoSink( target_path=target_many, video_info=info, codec="MJPG" ) as sink: for i in range(10): sink.write_frame(frame=np.full((64, 64, 3), i * 20, dtype=np.uint8)) assert Path(target_many).stat().st_size > Path(target_one).stat().st_size