refactor: improve tests

Author: Federico Pozzi

test/test_prototype_transforms.py

@@ -191,54 +191,56 @@ def test_convert_image_color_space(self, transform, input):
         transform(input)
 
 
+@pytest.mark.parametrize("p", [0.0, 1.0])
 class TestRandomHorizontalFlip:
-    def input_tensor(self, dtype=torch.float32):
-        return torch.tensor([[[0, 1], [0, 1]], [[1, 0], [1, 0]]], dtype=dtype)
+    def input_expected_image_tensor(self, p, dtype=torch.float32):
+        input = torch.tensor([[[0, 1], [0, 1]], [[1, 0], [1, 0]]], dtype=dtype)
+        expected = torch.tensor([[[1, 0], [1, 0]], [[0, 1], [0, 1]]], dtype=dtype)
 
-    def expected_tensor(self, dtype=torch.float32):
-        return torch.tensor([[[1, 0], [1, 0]], [[0, 1], [0, 1]]], dtype=dtype)
+        if p == 1.0:
+            return input, expected
+        return input, input
 
-    @pytest.mark.parametrize("p", [0.0, 1.0], ids=["p=0", "p=1"])
     def test_simple_tensor(self, p):
-        input = self.input_tensor()
+        input, expected = self.input_expected_image_tensor(p)
+        transform = transforms.RandomHorizontalFlip(p=p)
 
-        actual = transforms.RandomHorizontalFlip(p=p)(input)
+        actual = transform(input)
 
-        expected = self.expected_tensor() if p == 1.0 else input
         assert_equal(expected, actual)
 
-    @pytest.mark.parametrize("p", [0.0, 1.0], ids=["p=0", "p=1"])
     def test_pil_image(self, p):
-        input = self.input_tensor(dtype=torch.uint8)
+        input, expected = self.input_expected_image_tensor(p, dtype=torch.uint8)
+        transform = transforms.RandomHorizontalFlip(p=p)
 
-        actual = transforms.RandomHorizontalFlip(p=p)(to_pil_image(input))
+        actual = transform(to_pil_image(input))
 
-        expected = self.expected_tensor(dtype=torch.uint8) if p == 1.0 else input
         assert_equal(expected, pil_to_tensor(actual))
 
-    @pytest.mark.parametrize("p", [0.0, 1.0], ids=["p=0", "p=1"])
     def test_features_image(self, p):
-        input = self.input_tensor()
+        input, expected = self.input_expected_image_tensor(p)
+        transform = transforms.RandomHorizontalFlip(p=p)
 
-        actual = transforms.RandomHorizontalFlip(p=p)(features.Image(input))
+        actual = transform(features.Image(input))
 
-        expected = self.expected_tensor() if p == 1.0 else input
         assert_equal(features.Image(expected), actual)
 
-    @pytest.mark.parametrize("p", [0.0, 1.0], ids=["p=0", "p=1"])
     def test_features_segmentation_mask(self, p):
-        input = features.SegmentationMask(self.input_tensor())
+        input, expected = self.input_expected_image_tensor(p)
+        transform = transforms.RandomHorizontalFlip(p=p)
 
-        actual = transforms.RandomHorizontalFlip(p=p)(input)
+        actual = transform(features.SegmentationMask(input))
 
-        expected = self.expected_tensor() if p == 1.0 else input
         assert_equal(features.SegmentationMask(expected), actual)
 
-    @pytest.mark.parametrize("p", [0.0, 1.0], ids=["p=0", "p=1"])
     def test_features_bounding_box(self, p):
         input = features.BoundingBox([0, 0, 5, 5], format=features.BoundingBoxFormat.XYXY, image_size=(10, 10))
+        transform = transforms.RandomHorizontalFlip(p=p)
 
-        actual = transforms.RandomHorizontalFlip(p=p)(input)
+        actual = transform(input)
 
-        expected = torch.tensor([5, 0, 10, 5]) if p == 1.0 else input
-        assert_equal(features.BoundingBox.new_like(input, expected), actual)
+        expected_image_tensor = torch.tensor([5, 0, 10, 5]) if p == 1.0 else input
+        expected = features.BoundingBox.new_like(input, data=expected_image_tensor)
+        assert_equal(expected, actual)
+        assert actual.format == expected.format
+        assert actual.image_size == expected.image_size