TST parameterize imputer test
@@ -65,12 +65,13 @@ def _check_statistics(X, X_true,
assert_ae(X_trans, X_true, err_msg=err_msg.format(True))
-def test_imputation_shape():
+@pytest.mark.parametrize("strategy",
+ ['mean', 'median', 'most_frequent', "constant"])
+def test_imputation_shape(strategy):
# Verify the shapes of the imputed matrix for different strategies.
X = np.random.randn(10, 2)
X[::2] = np.nan
- for strategy in ['mean', 'median', 'most_frequent', "constant"]:
imputer = SimpleImputer(strategy=strategy)
X_imputed = imputer.fit_transform(sparse.csr_matrix(X))
assert X_imputed.shape == (10, 2)
