def test_extracts_more_than_one_col(self):
ce = ColsExtractor([0, 2])
features = ce.transform(self.X)
expected = [
[0, 2],
[3, 5]
]
self.assertTrue(np.array_equal(expected, features))
def test_extracts_one_col(self):
ce = ColsExtractor([1])
self.assertEqual([[1], [4]], ce.transform(self.X))
def test_extracts_and_flattens(self):
ce = ColsExtractor([0], flat=True)
features = ce.transform(self.X)
expected = [0, 3]
self.assertTrue(np.array_equal(expected, features))