TST refactor test_numeric_stability (#14221)

thomasjpfan · qinhanmin2014 · commit be4accf79d92 · 2019-06-30T10:28:29.000+08:00
diff --git a/sklearn/preprocessing/tests/test_discretization.py b/sklearn/preprocessing/tests/test_discretization.py
@@ -135,15 +135,15 @@ def test_transform_1d_behavior():
     assert_raises(ValueError, est.transform, X)
 
 
-def test_numeric_stability():
+@pytest.mark.parametrize('i', range(1, 9))
+def test_numeric_stability(i):
     X_init = np.array([2., 4., 6., 8., 10.]).reshape(-1, 1)
     Xt_expected = np.array([0, 0, 1, 1, 1]).reshape(-1, 1)
 
     # Test up to discretizing nano units
-    for i in range(1, 9):
-        X = X_init / 10**i
-        Xt = KBinsDiscretizer(n_bins=2, encode='ordinal').fit_transform(X)
-        assert_array_equal(Xt_expected, Xt)
+    X = X_init / 10**i
+    Xt = KBinsDiscretizer(n_bins=2, encode='ordinal').fit_transform(X)
+    assert_array_equal(Xt_expected, Xt)
 
 
 def test_invalid_encode_option():
