refactor to use pytest.mark.parameterize

Author: eendebakpt

numpy/_core/tests/test_numeric.py

@@ -1710,29 +1710,22 @@ def test_sparse(self):
             assert_equal(np.nonzero(c)[0],
                          np.concatenate((np.arange(10 + i, 20 + i), [20 + i*2])))
 
-    def test_nonzero_dtypes(self):
+    @pytest.mark.parametrize('dtype', [np.float32, np.float64])
+    def test_nonzero_float_dtypes(self, dtype):
         rng = np.random.default_rng(seed = 10)
-        zero_indices = np.arange(50)
-
-        # test for different dtypes
-        types = [bool, np.float32, np.float64]
-        sample = ((2**33)*rng.normal(size=100))
-        for dtype in types:
-            x = sample.astype(dtype)
-            rng.shuffle(zero_indices)
-            x[zero_indices] = 0
-            idxs = np.nonzero(x)[0]
-            assert_equal(np.array_equal(np.where(x != 0)[0], idxs), True)
-
-        integer_types = [np.int8, np.int16, np.int32, np.int64,
-                         np.uint8, np.uint16, np.uint32, np.uint64]
-        sample = rng.integers(0, 255, size=100)
-        for dtype in integer_types:
-            x = sample.astype(dtype)
-            rng.shuffle(zero_indices)
-            x[zero_indices] = 0
-            idxs = np.nonzero(x)[0]
-            assert_equal(np.array_equal(np.where(x != 0)[0], idxs), True)
+        x = ((2**33)*rng.normal(size=100)).astype(dtype)
+        x[rng.choice(50, size=100)] = 0
+        idxs = np.nonzero(x)[0]
+        assert_equal(np.array_equal(np.where(x != 0)[0], idxs), True)
+
+    @pytest.mark.parametrize('dtype', [bool, np.int8, np.int16, np.int32, np.int64,
+                                       np.uint8, np.uint16, np.uint32, np.uint64])
+    def test_nonzero_integer_dtypes(self, dtype):
+        rng = np.random.default_rng(seed = 10)
+        x = rng.integers(0, 255, size=100).astype(dtype)
+        x[rng.choice(50, size=100)] = 0
+        idxs = np.nonzero(x)[0]
+        assert_equal(np.array_equal(np.where(x != 0)[0], idxs), True)
 
 
     def test_return_type(self):