TST add test for sparse matrices

Author: glemaitre

imblearn/tests/test_misc.py

@@ -3,6 +3,10 @@
 # Authors: Guillaume Lemaitre <[email protected]>
 # License: MIT
 
+import pytest
+
+from scipy import sparse
+
 from sklearn.datasets import load_iris
 from sklearn.utils.testing import assert_array_equal
 
@@ -15,25 +19,40 @@
                       random_state=0)
 
 
-def test_function_sampler_identity():
+@pytest.mark.parametrize("X,y", [(X, y),
+                                 (sparse.csr_matrix(X), y),
+                                 (sparse.csc_matrix(X), y)])
+def test_function_sampler_identity(X, y):
     sampler = FunctionSampler()
     X_res, y_res = sampler.fit_sample(X, y)
+    if sparse.issparse(X):
+        X = X.toarray()
+        X_res = X_res.toarray()
     assert_array_equal(X_res, X)
     assert_array_equal(y_res, y)
 
 
-def test_function_sampler_func():
+@pytest.mark.parametrize("X,y", [(X, y),
+                                 (sparse.csr_matrix(X), y),
+                                 (sparse.csc_matrix(X), y)])
+def test_function_sampler_func(X, y):
 
     def func(X, y):
         return X[:10], y[:10]
 
     sampler = FunctionSampler(func=func)
     X_res, y_res = sampler.fit_sample(X, y)
+    if sparse.issparse(X):
+        X = X.toarray()
+        X_res = X_res.toarray()
     assert_array_equal(X_res, X[:10])
     assert_array_equal(y_res, y[:10])
 
 
-def test_function_sampler_func_kwargs():
+@pytest.mark.parametrize("X,y", [(X, y),
+                                 (sparse.csr_matrix(X), y),
+                                 (sparse.csc_matrix(X), y)])
+def test_function_sampler_func_kwargs(X, y):
 
     def func(X, y, ratio, random_state):
         rus = RandomUnderSampler(ratio=ratio, random_state=random_state)
@@ -43,5 +62,9 @@ def func(X, y, ratio, random_state):
                                                   'random_state': 0})
     X_res, y_res = sampler.fit_sample(X, y)
     X_res_2, y_res_2 = RandomUnderSampler(random_state=0).fit_sample(X, y)
+    if sparse.issparse(X):
+        X = X.toarray()
+        X_res = X_res.toarray()
+        X_res_2 = X_res_2.toarray()
     assert_array_equal(X_res, X_res_2)
     assert_array_equal(y_res, y_res_2)