def test_get_labelled_data_frame(self):
dataset = DataSet([[1, 2], [3, 4], [5, 6], [7, 8]],
labels=pd.Series(["b", "b", "b", "a"]))
df = dataset.get_labelled_data_frame()
# TODO: non-numeric values in DataFrame matcher
expected = [[1, 2, "b"], [3, 4, "b"], [5, 6, "b"], [7, 8, "a"]]
for i in range(len(expected)):
self.assertTrue(df.ix[i].tolist(), expected[i])
def test_get_labelled_data_frame_no_labels(self):
dataset = DataSet([[1, 2], [3, 4], [5, 6], [7, 8]])
df = dataset.get_labelled_data_frame()
assert_that(df, equals_dataframe([[1, 2], [3, 4], [5, 6], [7, 8]]))
