183 lines · 5.5 KB · Python
RawDownload
from fractions import Fractionfrom unittest.mock import Mock, call import avimport numpy as npimport pytest from manim._config.video_encoder import VideoEncoderSpecfrom manim.scene.video_segment_encoder import VideoSegmentEncoder  def _spec(*, width=4, height=2):    return VideoEncoderSpec(        container_format="mp4",        codec="libx264",        pixel_format="yuv420p",        width=width,        height=height,        frame_rate=Fraction(30, 1),        options=(("crf", "23"),),    )  def _detached_encoder(tmp_path, *, stream=None, container=None):    encoder = object.__new__(VideoSegmentEncoder)    encoder.target = tmp_path / "segment.mp4"    encoder.spec = _spec()    encoder._next_pts = 0    encoder._closed = False    encoder._stream = Mock() if stream is None else stream    encoder._container = Mock() if container is None else container    return encoder  def _frame():    return np.full((2, 4, 4), 127, dtype=np.uint8)  def test_write_frame_owns_repeat_and_segment_local_pts(tmp_path):    packet = object()    encoded_frames = []    stream = Mock()     def encode(frame):        encoded_frames.append(frame)        return [packet]     stream.encode.side_effect = encode    container = Mock()    encoder = _detached_encoder(tmp_path, stream=stream, container=container)     encoder.write_frame(_frame(), repeat=3)     assert [frame.pts for frame in encoded_frames] == [0, 1, 2]    assert [frame.time_base for frame in encoded_frames] == [Fraction(1, 30)] * 3    assert encoder._next_pts == 3    assert container.mux.call_args_list == [call(packet)] * 3  @pytest.mark.parametrize(    ("frame", "repeat", "exception", "message"),    [        (np.zeros((3, 4, 4), dtype=np.uint8), 1, ValueError, "shape"),        (np.zeros((2, 4, 3), dtype=np.uint8), 1, ValueError, "shape"),        (np.zeros((2, 4, 4), dtype=np.float32), 1, TypeError, "uint8"),        (np.zeros((2, 4, 4), dtype=np.uint8)[:, ::-1], 1, ValueError, "C-contiguous"),        (np.zeros((2, 4, 4), dtype=np.uint8), 0, ValueError, "positive integer"),        (np.zeros((2, 4, 4), dtype=np.uint8), True, ValueError, "positive integer"),    ],)def test_write_frame_validates_boundary(    tmp_path,    frame,    repeat,    exception,    message,):    encoder = _detached_encoder(tmp_path)     with pytest.raises(exception, match=message):        encoder.write_frame(frame, repeat=repeat)     encoder._stream.encode.assert_not_called()  def test_write_frame_rejects_closed_encoder(tmp_path):    encoder = _detached_encoder(tmp_path)    encoder._closed = True     with pytest.raises(RuntimeError, match="is closed"):        encoder.write_frame(_frame())  def test_encode_failure_has_target_profile_and_original_cause(tmp_path):    expected_exception = RuntimeError("codec exploded")    stream = Mock()    stream.encode.side_effect = expected_exception    encoder = _detached_encoder(tmp_path, stream=stream)     with pytest.raises(        RuntimeError, match=r"segment\.mp4.*mp4/libx264/yuv420p"    ) as exc_info:        encoder.write_frame(_frame())     assert exc_info.value.__cause__ is expected_exception  def test_finish_flushes_closes_and_is_idempotent(tmp_path):    packet = object()    stream = Mock()    stream.encode.return_value = [packet]    container = Mock()    encoder = _detached_encoder(tmp_path, stream=stream, container=container)     encoder.finish()    encoder.finish()     stream.encode.assert_called_once_with()    container.mux.assert_called_once_with(packet)    container.close.assert_called_once_with()  def test_finish_preserves_flush_failure_when_close_also_fails(tmp_path):    flush_exception = RuntimeError("flush failed")    close_exception = RuntimeError("close failed")    stream = Mock()    stream.encode.side_effect = flush_exception    container = Mock()    container.close.side_effect = close_exception    encoder = _detached_encoder(tmp_path, stream=stream, container=container)     with pytest.raises(RuntimeError, match="Failed to finish") as exc_info:        encoder.finish()     assert exc_info.value.__cause__ is flush_exception    container.close.assert_called_once_with()  def test_finish_closes_before_propagating_control_flow_exception(tmp_path):    stream = Mock()    stream.encode.side_effect = KeyboardInterrupt    container = Mock()    encoder = _detached_encoder(tmp_path, stream=stream, container=container)     with pytest.raises(KeyboardInterrupt):        encoder.finish()     container.close.assert_called_once_with()  def test_abort_closes_removes_target_and_is_idempotent(tmp_path):    encoder = _detached_encoder(tmp_path)    encoder.target.write_bytes(b"incomplete")     encoder.abort()    encoder.abort()     encoder._container.close.assert_called_once_with()    assert not encoder.target.exists()  def test_abort_removes_target_before_propagating_control_flow_exception(tmp_path):    container = Mock()    container.close.side_effect = KeyboardInterrupt    encoder = _detached_encoder(tmp_path, container=container)    encoder.target.write_bytes(b"incomplete")     with pytest.raises(KeyboardInterrupt):        encoder.abort()     assert not encoder.target.exists()  def test_open_failure_has_target_profile_and_original_cause(tmp_path, monkeypatch):    expected_exception = RuntimeError("open failed")    monkeypatch.setattr(av, "open", Mock(side_effect=expected_exception))     with pytest.raises(        RuntimeError, match=r"segment\.mp4.*mp4/libx264/yuv420p"    ) as exc_info:        VideoSegmentEncoder(target=tmp_path / "segment.mp4", spec=_spec())     assert exc_info.value.__cause__ is expected_exception