Improve test assertions to show actual values on failure (#1616)

Replace assertion patterns that don't show values on failure with NumPy
and PyTorch testing utilities.

Migrate from `self.assertTrue(np.all(...))` and
`self.assertTrue(torch.allclose(...))` to
`np.testing.assert_array_equal`, `np.testing.assert_array_compare`,
`np.testing.assert_allclose`, and `torch.testing.assert_close`. These
testing utilities automatically include summaries of actual vs expected
values in their error messages, making test failures much easier to
debug.

Author: mthrok

tests/cuda/buffer_transfer_test.py

@@ -69,7 +69,7 @@ def test_transfer_buffer_to_cuda(self, media_type: str) -> None:
         self.assertTrue(cuda_tensor.is_cuda)
         self.assertEqual(cuda_tensor.device, torch.device(f"cuda:{DEFAULT_CUDA}"))
 
-        self.assertTrue(torch.allclose(cpu_tensor, cuda_tensor.cpu()))
+        torch.testing.assert_close(cpu_tensor, cuda_tensor.cpu())
 
     @parameterized.expand(
         [
@@ -128,7 +128,7 @@ def test() -> None:
             cuda_tensor = spdl.io.to_torch(cuda_buffer)
             self.assertTrue(cuda_tensor.is_cuda)
             self.assertEqual(cuda_tensor.device, torch.device(f"cuda:{DEFAULT_CUDA}"))
-            self.assertTrue(torch.allclose(cpu_tensor, cuda_tensor.cpu()))
+            torch.testing.assert_close(cpu_tensor, cuda_tensor.cpu())
 
             print("Asserting deleter was not yet called")
             self.assertFalse(deleter_called)
@@ -161,7 +161,7 @@ def test(array) -> None:
             self.assertEqual(tensor.dtype, dtypes[array.dtype])
             self.assertEqual(tensor.shape, torch.Size(array.shape))
             self.assertEqual(tensor.device, torch.device(device))
-            self.assertTrue(torch.allclose(tensor, torch.from_numpy(array).to(device)))
+            torch.testing.assert_close(tensor, torch.from_numpy(array).to(device))
 
         for dtype in [np.uint8, np.int32, np.int64]:
             max_val = np.iinfo(dtype).max
@@ -182,7 +182,7 @@ def test(cpu_tensor) -> None:
             self.assertEqual(cuda_tensor.dtype, cpu_tensor.dtype)
             self.assertEqual(cuda_tensor.shape, cpu_tensor.shape)
             self.assertEqual(cuda_tensor.device, device)
-            self.assertTrue(torch.allclose(cuda_tensor, cpu_tensor.to(device)))
+            torch.testing.assert_close(cuda_tensor, cpu_tensor.to(device))
 
         for dtype in [np.uint8, np.int32, np.int64]:
             max_val = np.iinfo(dtype).max
@@ -205,7 +205,7 @@ def test_array_transfer_non_contiguous_torch(self) -> None:
         self.assertEqual(cuda_tensor.dtype, cpu_tensor.dtype)
         self.assertEqual(cuda_tensor.shape, cpu_tensor.shape)
         self.assertEqual(cuda_tensor.device, device)
-        self.assertTrue(torch.allclose(cuda_tensor, cpu_tensor.to(device)))
+        torch.testing.assert_close(cuda_tensor, cpu_tensor.to(device))
 
     def test_array_transfer_non_contiguous_numpy(self) -> None:
         """passing noncontiguous array/tensor to transfer_buffer works"""
@@ -224,7 +224,7 @@ def test_array_transfer_non_contiguous_numpy(self) -> None:
         self.assertEqual(tensor.dtype, torch.int64)
         self.assertEqual(tensor.shape, torch.Size(arr.shape))
         self.assertEqual(tensor.device, torch.device(device))
-        self.assertTrue(torch.allclose(tensor, torch.from_numpy(arr).to(device)))
+        torch.testing.assert_close(tensor, torch.from_numpy(arr).to(device))
 
     def test_array_transfer_smoke_test(self) -> None:
         """smoke test for transferring multiple arrays concurrently"""
@@ -253,4 +253,4 @@ def test_transfer_cpu(self) -> None:
             self.assertEqual(ref.dtype, cpu_tensor.dtype)
             self.assertEqual(ref.shape, cpu_tensor.shape)
             self.assertEqual(ref.device, cpu_tensor.device)
-            self.assertTrue(torch.allclose(ref, cpu_tensor))
+            torch.testing.assert_close(ref, cpu_tensor)

tests/cuda/nvjpeg_decode_test.py

@@ -44,9 +44,9 @@ def _test(data, pix_fmt):
         rgb_tensor = _test(sample.path, "rgb")
         bgr_tensor = _test(sample.path, "bgr")
 
-        self.assertTrue(torch.equal(rgb_tensor[0], bgr_tensor[2]))
-        self.assertTrue(torch.equal(rgb_tensor[1], bgr_tensor[1]))
-        self.assertTrue(torch.equal(rgb_tensor[2], bgr_tensor[0]))
+        torch.testing.assert_close(rgb_tensor[0], bgr_tensor[2])
+        torch.testing.assert_close(rgb_tensor[1], bgr_tensor[1])
+        torch.testing.assert_close(rgb_tensor[2], bgr_tensor[0])
 
     def test_decode_rubbish(self) -> None:
         """When decoding fails, it should raise an error instead of segfault then,

tests/cuda/pin_memory_test.py

@@ -88,7 +88,7 @@ def test_pin_memory_convert_array(self) -> None:
         self.assertEqual(tensor.shape, (10,))
         self.assertEqual(tensor.dtype, torch.int64)
         self.assertEqual(tensor.device, torch.device("cuda:0"))
-        self.assertTrue((vals == tensor.cpu().numpy()).all())
+        np.testing.assert_array_equal(tensor.cpu().numpy(), vals, strict=True)
 
     def test_pin_memory_convert_array_invalid_size(self) -> None:
         """convert_array fails if storage is small."""

tests/cuda/transfer_test.py

@@ -20,4 +20,4 @@ def test_gpu_transfer(self) -> None:
         cuda = transfer_tensor(ref)
         print(cuda)
         self.assertEqual(cuda.device.type, "cuda")
-        self.assertTrue(torch.equal(cuda, ref.cuda()))
+        torch.testing.assert_close(cuda, ref.cuda())

tests/dataloader/sampler_test.py

@@ -462,7 +462,7 @@ def test_weighted_sampling(self) -> None:
         print(f"{ref=}")
         print(f"{hyp=}")
 
-        self.assertTrue(np.allclose(hyp, ref, atol=1e-3))
+        np.testing.assert_allclose(hyp, ref, rtol=1e-5, atol=1e-3)
 
 
 class TestDistributedSamplerEmbedShuffle(unittest.TestCase):

tests/io/array_test.py

@@ -48,20 +48,20 @@ def test_load_npy_integral(
 
         data = _dump_npy(ref)
         recon = spdl.io.load_npy(data)
-        self.assertTrue(np.array_equal(recon, ref))
+        np.testing.assert_array_equal(recon, ref, strict=True)
 
         # Use bytearray to check if the change to the original is refrected to the recon
         # (which means that the recon is referring to the original, no copy)
         data = bytearray(data)
         print(f"{id(data)=}")
         print(f"{id(recon.data.obj)=}")
         recon = spdl.io.load_npy(data)
-        self.assertTrue(np.array_equal(recon, ref))
+        np.testing.assert_array_equal(recon, ref, strict=True)
 
         self.assertTrue(np.any(recon))
         # Fill zeros. The header is cleared too, but it's already parsed, so not an issue.
         data[:] = b"\x00" * len(data)
-        self.assertFalse(np.any(recon))
+        np.testing.assert_array_equal(recon, 0)
 
     @parameterized.expand(
         [
@@ -81,20 +81,20 @@ def test_load_npy_float(self, dtype: type[np.floating]) -> None:
 
         data = _dump_npy(ref)
         recon = spdl.io.load_npy(data)
-        self.assertTrue(np.array_equal(recon, ref))
+        np.testing.assert_array_equal(recon, ref, strict=True)
 
         # Use bytearray to check if the change to the original is refrected to the recon
         # (which means that the recon is referring to the original, no copy)
         data = bytearray(data)
         print(f"{id(data)=}")
         print(f"{id(recon.data.obj)=}")
         recon = spdl.io.load_npy(data)
-        self.assertTrue(np.array_equal(recon, ref))
+        np.testing.assert_array_equal(recon, ref, strict=True)
 
         self.assertTrue(np.any(recon))
         # Fill zeros. The header is cleared too, but it's already parsed, so not an issue.
         data[:] = b"\x00" * len(data)
-        self.assertFalse(np.any(recon))
+        np.testing.assert_array_equal(recon, 0)
 
 
 ##############################################################################

tests/io/async_test.py

@@ -91,14 +91,14 @@ def _decode(num_frames=None):
         arr1 = _decode(num_frames=num_frames)
         self.assertEqual(arr1.dtype, np.int16)
         self.assertEqual(arr1.shape, (num_frames, 1))
-        self.assertTrue(np.all(arr1 == arr0[:num_frames]))
+        np.testing.assert_array_equal(arr1, arr0[:num_frames], strict=True)
 
         num_frames = 32000
         arr2 = _decode(num_frames=num_frames)
         self.assertEqual(arr2.dtype, np.int16)
         self.assertEqual(arr2.shape, (num_frames, 1))
-        self.assertTrue(np.all(arr2[:16000] == arr0))
-        self.assertTrue(np.all(arr2[16000:] == 0))
+        np.testing.assert_array_equal(arr2[:16000], arr0, strict=True)
+        np.testing.assert_array_equal(arr2[16000:], 0)
 
     def test_decode_audio_many_channels_6(self) -> None:
         """Can decode audio with more than 6 channels.
@@ -229,28 +229,30 @@ def _decode(pix_fmt="rgb24", **kwargs):
         arr1 = _decode(num_frames=num_frames)
         self.assertEqual(arr1.dtype, np.uint8)
         self.assertEqual(arr1.shape, (num_frames, 240, 320, 3))
-        self.assertTrue(np.all(arr1 == arr0[:num_frames]))
+        np.testing.assert_array_equal(arr1, arr0[:num_frames], strict=True)
 
         num_frames = 100
         arr2 = _decode(num_frames=num_frames)
         self.assertEqual(arr2.dtype, np.uint8)
         self.assertEqual(arr2.shape, (num_frames, 240, 320, 3))
-        self.assertTrue(np.all(arr2[:50] == arr0))
-        self.assertTrue(np.all(arr2[50:] == arr2[50]))
+        np.testing.assert_array_equal(arr2[:50], arr0, strict=True)
+        np.testing.assert_array_equal(
+            arr2[50:], np.broadcast_to(arr2[50], arr2[50:].shape), strict=True
+        )
 
         num_frames = 100
         arr2 = _decode(num_frames=num_frames, pad_mode="black")
         self.assertEqual(arr2.dtype, np.uint8)
         self.assertEqual(arr2.shape, (num_frames, 240, 320, 3))
-        self.assertTrue(np.all(arr2[:50] == arr0))
-        self.assertTrue(np.all(arr2[50:] == 0))
+        np.testing.assert_array_equal(arr2[:50], arr0, strict=True)
+        np.testing.assert_array_equal(arr2[50:], 0)
 
         num_frames = 100
         arr2 = _decode(num_frames=num_frames, pad_mode="white")
         self.assertEqual(arr2.dtype, np.uint8)
         self.assertEqual(arr2.shape, (num_frames, 240, 320, 3))
-        self.assertTrue(np.all(arr2[:50] == arr0))
-        self.assertTrue(np.all(arr2[50:] == 255))
+        np.testing.assert_array_equal(arr2[:50], arr0, strict=True)
+        np.testing.assert_array_equal(arr2[50:], 255)
 
     def test_decode_video_frame_rate_pts(self) -> None:
         """Applying frame rate outputs correct PTS."""
@@ -267,7 +269,7 @@ def test_decode_video_frame_rate_pts(self) -> None:
         pts = frames.get_timestamps()
         print(pts_ref, pts)
 
-        self.assertTrue(np.all(pts_ref[::2] == pts))
+        np.testing.assert_array_equal(pts_ref[::2], pts, strict=True)
 
 
 class TestConvertFrames(unittest.TestCase):

tests/io/frames_clone_test.py

@@ -56,7 +56,7 @@ def test_clone_frames(self, media_type: str) -> None:
         array1 = _load_from_frames(frames1)
         array2 = _load_from_frames(frames2)
 
-        self.assertTrue(np.all(array1 == array2))
+        np.testing.assert_array_equal(array1, array2, strict=True)
 
     @parameterized.expand(
         [
@@ -113,4 +113,4 @@ def test_clone_frames_multi(self, media_type: str) -> None:
         arrays = [_load_from_frames(c) for c in clones]
 
         for i in range(N):
-            self.assertTrue(np.all(array == arrays[i]))
+            np.testing.assert_array_equal(array, arrays[i], strict=True)

tests/io/image_decoding_test.py

@@ -6,6 +6,7 @@
 
 # pyre-unsafe
 
+import operator
 import unittest
 
 import numpy as np
@@ -21,6 +22,18 @@ def _load_image(src, filter_desc="format=pix_fmts=rgb24"):
     return spdl.io.to_numpy(spdl.io.load_image(src, filter_desc=filter_desc))
 
 
+def _assert_all_ge(x, val) -> None:
+    np.testing.assert_array_compare(
+        operator.ge, x, val, header=f"Arrays are not all >= {val}"
+    )
+
+
+def _assert_all_le(x, val) -> None:
+    np.testing.assert_array_compare(
+        operator.le, x, val, header=f"Arrays are not all <= {val}"
+    )
+
+
 class TestDecodeImage(unittest.TestCase):
     def test_decode_image_gray16_native(self) -> None:
         """Can decode gray16 PNG image (16be) as-is"""
@@ -44,8 +57,8 @@ def test_decode_image_gray16_native(self) -> None:
 
         ffmpeg8 = get_ffmpeg_versions()["libavutil"][0] >= 60
 
-        self.assertTrue(np.all(hyp[..., :32] == 65535 if ffmpeg8 else 65022))
-        self.assertTrue(np.all(hyp[..., 32:] == 0 if ffmpeg8 else 256))
+        np.testing.assert_array_equal(hyp[..., :32], 65535 if ffmpeg8 else 65022)
+        np.testing.assert_array_equal(hyp[..., 32:], 0 if ffmpeg8 else 256)
 
     def test_decode_image_16be_rgb24(self) -> None:
         """Can decode gray16 PNG image (16be) as rgb24"""
@@ -67,8 +80,8 @@ def test_decode_image_16be_rgb24(self) -> None:
         ref = load_ref_image(sample.path, shape)
         np.testing.assert_array_equal(hyp, ref, strict=True)
 
-        self.assertTrue(np.all(hyp[:, :32, :] == 255))
-        self.assertTrue(np.all(hyp[:, 32:, :] == 0))
+        np.testing.assert_array_equal(hyp[:, :32, :], 255)
+        np.testing.assert_array_equal(hyp[:, 32:, :], 0)
 
     def test_decode_image_yuvj422_native(self) -> None:
         """Can decode yuvj422p JPEG image as-is."""
@@ -153,17 +166,17 @@ def test_decode_image_yuvj420p_as_rgb24_edge_values(self) -> None:
 
         red, green, blue = hyp[:, :width], hyp[:, width:2*width], hyp[:, 2*width:]
 
-        self.assertTrue(np.all(red[..., 0] >= 254))
-        self.assertTrue(np.all(red[..., 1] <= 1))
-        self.assertTrue(np.all(red[..., 2] == 0))
+        _assert_all_ge(red[..., 0], 254)
+        _assert_all_le(red[..., 1], 1)
+        np.testing.assert_array_equal(red[..., 2], 0)
 
-        self.assertTrue(np.all(green[..., 0] == 0))
-        self.assertTrue(np.all(green[..., 1] >= 253))
-        self.assertTrue(np.all(green[..., 2] <= 1))
+        np.testing.assert_array_equal(green[..., 0], 0)
+        _assert_all_ge(green[..., 1], 253)
+        _assert_all_le(green[..., 2], 1)
 
-        self.assertTrue(np.all(blue[..., 0] <= 1))
-        self.assertTrue(np.all(blue[..., 1] <= 1))
-        self.assertTrue(np.all(blue[..., 2] >= 254))
+        _assert_all_le(blue[..., 0], 1)
+        _assert_all_le(blue[..., 1], 1)
+        _assert_all_ge(blue[..., 2], 254)
 
     def test_decode_image_yuvj444p_native(self) -> None:
         """Can decode yuvj444p JPEG image as-is."""
@@ -335,19 +348,19 @@ def test_load_image_batch_native_edge_values(self) -> None:
 
         left, middle, right = arr[..., :w, :], arr[..., w:-w, :], arr[..., -w:, :]
         # Red
-        self.assertTrue(np.all(left[..., 0] >= 252))
-        self.assertTrue(np.all(left[..., 1] == 0))
-        self.assertTrue(np.all(left[..., 2] == 0))
+        _assert_all_ge(left[..., 0], 252)
+        np.testing.assert_array_equal(left[..., 1], 0)
+        np.testing.assert_array_equal(left[..., 2], 0)
 
         # Green
-        self.assertTrue(np.all(middle[..., 0] == 0))
-        self.assertTrue(np.all(middle[..., 1] >= 253))
-        self.assertTrue(np.all(middle[..., 2] == 0))
+        np.testing.assert_array_equal(middle[..., 0], 0)
+        _assert_all_ge(middle[..., 1], 253)
+        np.testing.assert_array_equal(middle[..., 2], 0)
 
         # Blue
-        self.assertTrue(np.all(right[..., 0] == 0))
-        self.assertTrue(np.all(right[..., 1] == 0))
-        self.assertTrue(np.all(right[..., 2] >= 253))
+        np.testing.assert_array_equal(right[..., 0], 0)
+        np.testing.assert_array_equal(right[..., 1], 0)
+        _assert_all_ge(right[..., 2], 253)
 
     def test_batch_decode_image_handle_failure(self) -> None:
         """load_image_batch dismisses failures when strict=False."""