iter

Author: glemaitre

imblearn/utils/tests/test_validation.py

@@ -9,32 +9,21 @@
 import pytest
 import numpy as np
 
-from sklearn.base import BaseEstimator
 from sklearn.neighbors._base import KNeighborsMixin
 from sklearn.neighbors import NearestNeighbors
 from sklearn.utils._testing import assert_array_equal
 
-from imblearn.utils.testing import warns
 from imblearn.utils import check_neighbors_object
 from imblearn.utils import check_sampling_strategy
 from imblearn.utils import check_target_type
+from imblearn.utils.testing import warns, _CustomNearestNeighbors
 from imblearn.utils._validation import ArraysTransformer
 from imblearn.utils._validation import _deprecate_positional_args
 
 multiclass_target = np.array([1] * 50 + [2] * 100 + [3] * 25)
 binary_target = np.array([1] * 25 + [0] * 100)
 
 
-class KNNLikeEstimator(BaseEstimator):
-    """A class exposing the same KNeighborsMixin API than KNeighborsClassifier."""
-
-    def kneighbors(self, X):
-        return np.ones((len(X), 1))
-
-    def kneighbors_graph(self, X):
-        return np.ones((len(X), 1))
-
-
 def test_check_neighbors_object():
     name = "n_neighbors"
     n_neighbors = 1
@@ -47,9 +36,9 @@ def test_check_neighbors_object():
     estimator = NearestNeighbors(n_neighbors=n_neighbors)
     estimator_cloned = check_neighbors_object(name, estimator)
     assert estimator.n_neighbors == estimator_cloned.n_neighbors
-    estimator = KNNLikeEstimator()
+    estimator = _CustomNearestNeighbors()
     estimator_cloned = check_neighbors_object(name, estimator)
-    assert isinstance(estimator_cloned, KNNLikeEstimator)
+    assert isinstance(estimator_cloned, _CustomNearestNeighbors)
     n_neighbors = "rnd"
     err_msg = (
         "n_neighbors must be an interger or an object compatible with the "