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)