Merge pull request #29 from ogrisel/enh-pr-5540

amueller · amueller · commit 082c7f1adcac · 2015-10-23T16:16:29.000+02:00
TST improved test for df param
diff --git a/sklearn/tests/test_base.py b/sklearn/tests/test_base.py
@@ -237,13 +237,9 @@ def test_score_sample_weight():
 
 
 def test_clone_pandas_dataframe():
-    class MockDataFrameWithEq(MockDataFrame):
-        """implemenets __eq__ operator to leverage the current test."""
-        def __eq__(self, other):
-            return self.array == other.array
 
     class DummyEstimator(BaseEstimator, TransformerMixin):
-        """This is a dummpy class for generating numerical features
+        """This is a dummy class for generating numerical features
 
         This feature extractor extracts numerical features from pandas data
         frame.
@@ -257,9 +253,9 @@ class DummyEstimator(BaseEstimator, TransformerMixin):
         Notes
         -----
         """
-        def __init__(self, df, toto):
+        def __init__(self, df=None, scalar_param=1):
             self.df = df
-            self.toto = toto
+            self.scalar_param = scalar_param
 
         def fit(self, X, y=None):
             pass
@@ -269,9 +265,10 @@ def transform(self, X, y=None):
 
     # build and clone estimator
     d = np.arange(10)
-    df = MockDataFrameWithEq(d)
-    e = DummyEstimator(df, toto=1)
+    df = MockDataFrame(d)
+    e = DummyEstimator(df, scalar_param=1)
     cloned_e = clone(e)
 
     # the test
-    assert_equal(e.toto, cloned_e.toto)
+    assert_true((e.df == cloned_e.df).values.all())
+    assert_equal(e.scalar_param, cloned_e.scalar_param)
diff --git a/sklearn/utils/mocking.py b/sklearn/utils/mocking.py
@@ -18,6 +18,7 @@ class MockDataFrame(object):
     # have shape an length but don't support indexing.
     def __init__(self, array):
         self.array = array
+        self.values = array
         self.shape = array.shape
         self.ndim = array.ndim
         # ugly hack to make iloc work.
@@ -32,6 +33,9 @@ def __array__(self):
         # method.
         return self.array
 
+    def __eq__(self, other):
+        return MockDataFrame(self.array == other.array)
+
 
 class CheckingClassifier(BaseEstimator, ClassifierMixin):
     """Dummy classifier to test pipelining and meta-estimators.
