Also check fields of AudioSamples

Author: NicolasHug

test/test_encoders.py

@@ -53,7 +53,7 @@ def validate_frames_properties(*, actual: Path, expected: Path):
     for frame_index, (d_actual, d_expected) in enumerate(
         zip(frames_actual, frames_expected)
     ):
-        for prop in d_actual:
+        for prop in d_expected:
             if prop == "pkt_pos":
                 continue  # TODO this probably matters
             assert (
@@ -66,7 +66,7 @@ class TestAudioEncoder:
     def decode(self, source) -> torch.Tensor:
         if isinstance(source, TestContainerFile):
             source = str(source.path)
-        return AudioDecoder(source).get_all_samples().data
+        return AudioDecoder(source).get_all_samples()
 
     def test_bad_input(self):
         with pytest.raises(ValueError, match="Expected samples to be a Tensor"):
@@ -108,12 +108,12 @@ def test_bad_input_parametrized(self, method, tmp_path):
             else dict(format="mp3")
         )
 
-        decoder = AudioEncoder(self.decode(NASA_AUDIO_MP3), sample_rate=10)
+        decoder = AudioEncoder(self.decode(NASA_AUDIO_MP3).data, sample_rate=10)
         with pytest.raises(RuntimeError, match="invalid sample rate=10"):
             getattr(decoder, method)(**valid_params)
 
         decoder = AudioEncoder(
-            self.decode(NASA_AUDIO_MP3), sample_rate=NASA_AUDIO_MP3.sample_rate
+            self.decode(NASA_AUDIO_MP3).data, sample_rate=NASA_AUDIO_MP3.sample_rate
         )
         with pytest.raises(RuntimeError, match="bit_rate=-1 must be >= 0"):
             getattr(decoder, method)(**valid_params, bit_rate=-1)
@@ -126,7 +126,7 @@ def test_bad_input_parametrized(self, method, tmp_path):
             getattr(decoder, method)(**valid_params)
 
         decoder = AudioEncoder(
-            self.decode(NASA_AUDIO_MP3), sample_rate=NASA_AUDIO_MP3.sample_rate
+            self.decode(NASA_AUDIO_MP3).data, sample_rate=NASA_AUDIO_MP3.sample_rate
         )
         for num_channels in (0, 3):
             with pytest.raises(
@@ -146,7 +146,7 @@ def test_round_trip(self, method, format, tmp_path):
             pytest.skip("Swresample with FFmpeg 4 doesn't work on wav files")
 
         asset = NASA_AUDIO_MP3
-        source_samples = self.decode(asset)
+        source_samples = self.decode(asset).data
 
         encoder = AudioEncoder(source_samples, sample_rate=asset.sample_rate)
 
@@ -161,7 +161,7 @@ def test_round_trip(self, method, format, tmp_path):
 
         rtol, atol = (0, 1e-4) if format == "wav" else (None, None)
         torch.testing.assert_close(
-            self.decode(encoded_source), source_samples, rtol=rtol, atol=atol
+            self.decode(encoded_source).data, source_samples, rtol=rtol, atol=atol
         )
 
     @pytest.mark.skipif(in_fbcode(), reason="TODO: enable ffmpeg CLI")
@@ -189,7 +189,7 @@ def test_against_cli(self, asset, bit_rate, num_channels, format, method, tmp_pa
             check=True,
         )
 
-        encoder = AudioEncoder(self.decode(asset), sample_rate=asset.sample_rate)
+        encoder = AudioEncoder(self.decode(asset).data, sample_rate=asset.sample_rate)
         params = dict(bit_rate=bit_rate, num_channels=num_channels)
         if method == "to_file":
             encoded_by_us = tmp_path / f"output.{format}"
@@ -207,13 +207,17 @@ def test_against_cli(self, asset, bit_rate, num_channels, format, method, tmp_pa
             rtol, atol = 0, 1e-3
         else:
             rtol, atol = None, None
-        # TODO should validate `.pts_seconds` and `duration_seconds` as well
+        samples_by_us = self.decode(encoded_by_us)
+        samples_by_ffmpeg = self.decode(encoded_by_ffmpeg)
         torch.testing.assert_close(
-            self.decode(encoded_by_us),
-            self.decode(encoded_by_ffmpeg),
+            samples_by_us.data,
+            samples_by_ffmpeg.data,
             rtol=rtol,
             atol=atol,
         )
+        assert samples_by_us.pts_seconds == samples_by_ffmpeg.pts_seconds
+        assert samples_by_us.duration_seconds == samples_by_ffmpeg.duration_seconds
+        assert samples_by_us.sample_rate == samples_by_ffmpeg.sample_rate
 
         if method == "to_file":
             validate_frames_properties(actual=encoded_by_us, expected=encoded_by_ffmpeg)
@@ -230,7 +234,7 @@ def test_to_tensor_against_to_file(
         if get_ffmpeg_major_version() == 4 and format == "wav":
             pytest.skip("Swresample with FFmpeg 4 doesn't work on wav files")
 
-        encoder = AudioEncoder(self.decode(asset), sample_rate=asset.sample_rate)
+        encoder = AudioEncoder(self.decode(asset).data, sample_rate=asset.sample_rate)
 
         params = dict(bit_rate=bit_rate, num_channels=num_channels)
         encoded_file = tmp_path / f"output.{format}"
@@ -240,7 +244,7 @@ def test_to_tensor_against_to_file(
         )
 
         torch.testing.assert_close(
-            self.decode(encoded_file), self.decode(encoded_tensor)
+            self.decode(encoded_file).data, self.decode(encoded_tensor).data
         )
 
     def test_encode_to_tensor_long_output(self):
@@ -256,7 +260,7 @@ def test_encode_to_tensor_long_output(self):
         INITIAL_TENSOR_SIZE = 10_000_000
         assert encoded_tensor.numel() > INITIAL_TENSOR_SIZE
 
-        torch.testing.assert_close(self.decode(encoded_tensor), samples)
+        torch.testing.assert_close(self.decode(encoded_tensor).data, samples)
 
     def test_contiguity(self):
         # Ensure that 2 waveforms with the same values are encoded in the same
@@ -313,4 +317,4 @@ def test_num_channels(
 
         if num_channels_output is None:
             num_channels_output = num_channels_input
-        assert self.decode(encoded_source).shape[0] == num_channels_output
+        assert self.decode(encoded_source).data.shape[0] == num_channels_output