def test_ColumnSelector__pick_three_columns(self): X = self.X selector = ColumnSelector(mapping={'X': slice(0, 3)}) self.assertTrue(selector.fit() is selector) assert_frame_equal(selector.predict(X)['X'], X)
def test_ColumnSelector__pick_one_column_last(self): X = self.X selector = ColumnSelector(mapping={'y': slice(2, 3)}) self.assertTrue(selector.fit() is selector) assert_frame_equal(selector.predict(X)['y'], X.loc[:, ["V3"]])
def test_ColumnSelector__pick_two_columns(self): X = self.X selector = ColumnSelector(mapping={'X': slice(0, 2)}) self.assertTrue(selector.fit() is selector) assert_frame_equal(selector.predict(X)['X'], X.loc[:, ["V1", "V2"]])
def test_ColumnSelector__mapping_is_None(self): X = self.X selector = ColumnSelector() self.assertTrue(selector.fit() is selector) assert_frame_equal(selector.predict(X)['predict'], X)