TST Extend tests for scipy.sparse/*array in sklearn/linear_model/tests/test_ransac (scikit-learn#27233)

Charlie-XIAO · REDVM · commit 72560c190eb9 · 2023-11-16T16:28:15.000+01:00
diff --git a/sklearn/linear_model/tests/test_ransac.py b/sklearn/linear_model/tests/test_ransac.py
@@ -1,7 +1,6 @@
 import numpy as np
 import pytest
 from numpy.testing import assert_array_almost_equal, assert_array_equal
-from scipy import sparse
 
 from sklearn.datasets import make_regression
 from sklearn.exceptions import ConvergenceWarning
@@ -14,6 +13,7 @@
 from sklearn.linear_model._ransac import _dynamic_max_trials
 from sklearn.utils import check_random_state
 from sklearn.utils._testing import assert_allclose
+from sklearn.utils.fixes import COO_CONTAINERS, CSC_CONTAINERS, CSR_CONTAINERS
 
 # Generate coordinates of line
 X = np.arange(-200, 200)
@@ -248,38 +248,11 @@ def is_data_valid(X, y):
     assert ransac_estimator.n_skips_invalid_model_ == 0
 
 
-def test_ransac_sparse_coo():
-    X_sparse = sparse.coo_matrix(X)
-
-    estimator = LinearRegression()
-    ransac_estimator = RANSACRegressor(
-        estimator, min_samples=2, residual_threshold=5, random_state=0
-    )
-    ransac_estimator.fit(X_sparse, y)
-
-    ref_inlier_mask = np.ones_like(ransac_estimator.inlier_mask_).astype(np.bool_)
-    ref_inlier_mask[outliers] = False
-
-    assert_array_equal(ransac_estimator.inlier_mask_, ref_inlier_mask)
-
-
-def test_ransac_sparse_csr():
-    X_sparse = sparse.csr_matrix(X)
-
-    estimator = LinearRegression()
-    ransac_estimator = RANSACRegressor(
-        estimator, min_samples=2, residual_threshold=5, random_state=0
-    )
-    ransac_estimator.fit(X_sparse, y)
-
-    ref_inlier_mask = np.ones_like(ransac_estimator.inlier_mask_).astype(np.bool_)
-    ref_inlier_mask[outliers] = False
-
-    assert_array_equal(ransac_estimator.inlier_mask_, ref_inlier_mask)
-
-
-def test_ransac_sparse_csc():
-    X_sparse = sparse.csc_matrix(X)
+@pytest.mark.parametrize(
+    "sparse_container", COO_CONTAINERS + CSR_CONTAINERS + CSC_CONTAINERS
+)
+def test_ransac_sparse(sparse_container):
+    X_sparse = sparse_container(X)
 
     estimator = LinearRegression()
     ransac_estimator = RANSACRegressor(
