Add test

Author: NicolasHug

src/torchcodec/decoders/_core/VideoDecoder.cpp

@@ -1324,7 +1324,8 @@ void VideoDecoder::convertAudioAVFrameToFrameOutputOnCPU(
 
   auto numSamples = avFrame->nb_samples; // per channel
   auto numChannels = getNumChannels(avFrame);
-  torch::Tensor outputData = torch::empty({numChannels, numSamples}, torch::kFloat32);
+  torch::Tensor outputData =
+      torch::empty({numChannels, numSamples}, torch::kFloat32);
 
   AVSampleFormat format = static_cast<AVSampleFormat>(avFrame->format);
   // TODO-AUDIO Implement all formats.

test/decoders/test_ops.py

@@ -695,7 +695,7 @@ def test_get_frames_by_pts_in_range_audio(self, range, asset):
             decoder, start_seconds=start_seconds, stop_seconds=stop_seconds
         )
 
-        assert_frames_equal(frames, reference_frames)
+        torch.testing.assert_close(frames, reference_frames)
 
     @pytest.mark.parametrize(
         "asset, expected_shape", ((NASA_AUDIO, (2, 1024)), (NASA_AUDIO_MP3, (2, 576)))
@@ -723,6 +723,46 @@ def test_decode_just_one_frame_at_boundaries(self, asset, expected_shape):
         )
         assert frames.shape == expected_shape
 
+    @pytest.mark.parametrize("asset", (NASA_AUDIO, NASA_AUDIO_MP3))
+    def test_multiple_calls(self, asset):
+
+        def decode_stateless(start_seconds, stop_seconds):
+            decoder = create_from_file(str(asset.path), seek_mode="approximate")
+            add_audio_stream(decoder)
+
+            return get_frames_by_pts_in_range_audio(
+                decoder, start_seconds=start_seconds, stop_seconds=stop_seconds
+            )
+
+        decoder = create_from_file(str(asset.path), seek_mode="approximate")
+        add_audio_stream(decoder)
+
+        start_seconds, stop_seconds = 0, 2
+        frames = get_frames_by_pts_in_range_audio(
+            decoder, start_seconds=start_seconds, stop_seconds=stop_seconds
+        )
+        torch.testing.assert_close(
+            frames, decode_stateless(start_seconds, stop_seconds)
+        )
+
+        start_seconds, stop_seconds = 3, 4
+        frames = get_frames_by_pts_in_range_audio(
+            decoder, start_seconds=start_seconds, stop_seconds=stop_seconds
+        )
+        torch.testing.assert_close(
+            frames, decode_stateless(start_seconds, stop_seconds)
+        )
+
+        # TODO-AUDIO
+        start_seconds, stop_seconds = 0, 2
+        frames = get_frames_by_pts_in_range_audio(
+            decoder, start_seconds=start_seconds, stop_seconds=stop_seconds
+        )
+        with pytest.raises(AssertionError):
+            torch.testing.assert_close(
+                frames, decode_stateless(start_seconds, stop_seconds)
+            )
+
 
 if __name__ == "__main__":
     pytest.main()

test/utils.py

@@ -26,25 +26,16 @@ def cpu_and_cuda():
     return ("cpu", pytest.param("cuda", marks=pytest.mark.needs_cuda))
 
 
-def assert_frames_equal(*args, **kwargs):
-    frame = args[0]
-    # This heuristic will work until we start returning uint8 audio frames...
-    if frame.dtype == torch.uint8:
-        return assert_video_frames_equal(*args, **kwargs)
-    else:
-        return assert_audio_frames_equal(*args, **kwargs)
-
-
-def assert_audio_frames_equal(*args, **kwargs):
-    torch.testing.assert_close(*args, **kwargs)
+def get_ffmpeg_major_version():
+    return int(get_ffmpeg_library_versions()["ffmpeg_version"].split(".")[0])
 
 
 # For use with decoded data frames. On CPU Linux, we expect exact, bit-for-bit
 # equality. On CUDA Linux, we expect a small tolerance.
 # On other platforms (e.g. MacOS), we also allow a small tolerance. FFmpeg does
 # not guarantee bit-for-bit equality across systems and architectures, so we
 # also cannot. We currently use Linux on x86_64 as our reference system.
-def assert_video_frames_equal(*args, **kwargs):
+def assert_frames_equal(*args, **kwargs):
     if sys.platform == "linux":
         if args[0].device.type == "cuda":
             atol = 2
@@ -82,9 +73,6 @@ def assert_tensor_close_on_at_least(actual_tensor, ref_tensor, *, percentage, at
         )
 
 
-def get_ffmpeg_major_version():
-    return int(get_ffmpeg_library_versions()["ffmpeg_version"].split(".")[0])
-
 
 def in_fbcode() -> bool:
     return os.environ.get("IN_FBCODE_TORCHCODEC") == "1"