From 54e995054d341fb445311326486d7d7a6d7727e5 Mon Sep 17 00:00:00 2001
From: Nicolas Hug <nicolashug@meta.com>
Date: Tue, 31 Oct 2023 05:02:33 -0700
Subject: [PATCH] [fbsync] port tests for transforms.Lambda (#8011)

Reviewed By: vmoens

Differential Revision: D50789100

fbshipit-source-id: 96b51f5e25569a943cbafe04f9afcd57e2f94959
---
 test/test_transforms_v2.py             | 35 --------------------------
 test/test_transforms_v2_consistency.py | 10 --------
 test/test_transforms_v2_refactored.py  | 18 +++++++++++++
 3 files changed, 18 insertions(+), 45 deletions(-)

diff --git a/test/test_transforms_v2.py b/test/test_transforms_v2.py
index 03f9e906675..61ab0d6a9ba 100644
--- a/test/test_transforms_v2.py
+++ b/test/test_transforms_v2.py
@@ -574,38 +574,3 @@ def test_sanitize_bounding_boxes_errors():
     with pytest.raises(ValueError, match="Number of boxes"):
         different_sizes = {"bbox": good_bbox, "labels": torch.arange(good_bbox.shape[0] + 3)}
         transforms.SanitizeBoundingBoxes()(different_sizes)
-
-
-class TestLambda:
-    inputs = pytest.mark.parametrize("input", [object(), torch.empty(()), np.empty(()), "string", 1, 0.0])
-
-    @inputs
-    def test_default(self, input):
-        was_applied = False
-
-        def was_applied_fn(input):
-            nonlocal was_applied
-            was_applied = True
-            return input
-
-        transform = transforms.Lambda(was_applied_fn)
-
-        transform(input)
-
-        assert was_applied
-
-    @inputs
-    def test_with_types(self, input):
-        was_applied = False
-
-        def was_applied_fn(input):
-            nonlocal was_applied
-            was_applied = True
-            return input
-
-        types = (torch.Tensor, np.ndarray)
-        transform = transforms.Lambda(was_applied_fn, *types)
-
-        transform(input)
-
-        assert was_applied is isinstance(input, types)
diff --git a/test/test_transforms_v2_consistency.py b/test/test_transforms_v2_consistency.py
index 397d42101ce..622ff4fe491 100644
--- a/test/test_transforms_v2_consistency.py
+++ b/test/test_transforms_v2_consistency.py
@@ -72,16 +72,6 @@ def __init__(
 LINEAR_TRANSFORMATION_MATRIX = torch.rand([LINEAR_TRANSFORMATION_MEAN.numel()] * 2)
 
 CONSISTENCY_CONFIGS = [
-    ConsistencyConfig(
-        v2_transforms.Lambda,
-        legacy_transforms.Lambda,
-        [
-            NotScriptableArgsKwargs(lambda image: image / 2),
-        ],
-        # Technically, this also supports PIL, but it is overkill to write a function here that supports tensor and PIL
-        # images given that the transform does nothing but call it anyway.
-        supports_pil=False,
-    ),
     ConsistencyConfig(
         v2_transforms.Compose,
         legacy_transforms.Compose,
diff --git a/test/test_transforms_v2_refactored.py b/test/test_transforms_v2_refactored.py
index b700b159ec5..8197e26eee7 100644
--- a/test/test_transforms_v2_refactored.py
+++ b/test/test_transforms_v2_refactored.py
@@ -5126,3 +5126,21 @@ def test_functional_and_transform(self, color_space, fn):
     def test_functional_error(self):
         with pytest.raises(TypeError, match="pic should be PIL Image"):
             F.pil_to_tensor(object())
+
+
+class TestLambda:
+    @pytest.mark.parametrize("input", [object(), torch.empty(()), np.empty(()), "string", 1, 0.0])
+    @pytest.mark.parametrize("types", [(), (torch.Tensor, np.ndarray)])
+    def test_transform(self, input, types):
+        was_applied = False
+
+        def was_applied_fn(input):
+            nonlocal was_applied
+            was_applied = True
+            return input
+
+        transform = transforms.Lambda(was_applied_fn, *types)
+        output = transform(input)
+
+        assert output is input
+        assert was_applied is (not types or isinstance(input, types))