roboflow--supervision/tests/utils/test_sinks.py

"""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