Moved tests out of run_in_graph_and_eager_mode in connected components. (#1412)

See #1328

Author: gabrieldemarmiesse

tensorflow_addons/image/connected_components_test.py

@@ -41,36 +41,38 @@
 )
 
 
-@test_utils.run_all_in_graph_and_eager_modes
-class ConnectedComponentsTest(tf.test.TestCase):
-    def testDisconnected(self):
-        arr = tf.cast(
-            [
-                [1, 0, 0, 1, 0, 0, 0, 0, 1],
-                [0, 1, 0, 0, 0, 1, 0, 1, 0],
-                [1, 0, 1, 0, 0, 0, 1, 0, 0],
-                [0, 0, 0, 0, 1, 0, 0, 0, 0],
-                [0, 0, 1, 0, 0, 0, 0, 0, 0],
-            ],
-            tf.bool,
-        )
-        expected = [
-            [1, 0, 0, 2, 0, 0, 0, 0, 3],
-            [0, 4, 0, 0, 0, 5, 0, 6, 0],
-            [7, 0, 8, 0, 0, 0, 9, 0, 0],
-            [0, 0, 0, 0, 10, 0, 0, 0, 0],
-            [0, 0, 11, 0, 0, 0, 0, 0, 0],
-        ]
-        self.assertAllEqual(self.evaluate(connected_components(arr)), expected)
+@pytest.mark.usefixtures("maybe_run_functions_eagerly")
+def test_disconnected():
+    arr = tf.cast(
+        [
+            [1, 0, 0, 1, 0, 0, 0, 0, 1],
+            [0, 1, 0, 0, 0, 1, 0, 1, 0],
+            [1, 0, 1, 0, 0, 0, 1, 0, 0],
+            [0, 0, 0, 0, 1, 0, 0, 0, 0],
+            [0, 0, 1, 0, 0, 0, 0, 0, 0],
+        ],
+        tf.bool,
+    )
+    expected = [
+        [1, 0, 0, 2, 0, 0, 0, 0, 3],
+        [0, 4, 0, 0, 0, 5, 0, 6, 0],
+        [7, 0, 8, 0, 0, 0, 9, 0, 0],
+        [0, 0, 0, 0, 10, 0, 0, 0, 0],
+        [0, 0, 11, 0, 0, 0, 0, 0, 0],
+    ]
+    np.testing.assert_equal(connected_components(arr).numpy(), expected)
 
-    def testSimple(self):
-        arr = [[0, 1, 0], [1, 1, 1], [0, 1, 0]]
 
-        # Single component with id 1.
-        self.assertAllEqual(
-            self.evaluate(connected_components(tf.cast(arr, tf.bool))), arr
-        )
+@pytest.mark.usefixtures("maybe_run_functions_eagerly")
+def test_simple():
+    arr = [[0, 1, 0], [1, 1, 1], [0, 1, 0]]
 
+    # Single component with id 1.
+    np.testing.assert_equal(connected_components(tf.cast(arr, tf.bool)).numpy(), arr)
+
+
+@test_utils.run_all_in_graph_and_eager_modes
+class ConnectedComponentsTest(tf.test.TestCase):
     def testSnake(self):
 
         # Single component with id 1.