def test_get_column(self):
lr = [["a", "b"], [1, 2], [3, 4], [5, 6]]
df = DataFrame(lr)
# test normal usage with string
result = df.get_column("b")
self.assertEqual(result, [2, 4, 6])
# test normal usage with integers
result = df.get_column(1)
self.assertEqual(result, [2, 4, 6])
# test invalid head
with self.assertRaises(KeyError) as cm:
result = df.get_column("c")
e = cm.exception
self.assertEqual(e.__str__(), "'There is no column head called c'"
) # '' apparently needed, because of how KeyError
def test_append_column(self):
# test correct use
df = DataFrame([["a", "b"], [1, 2], [4, 5], [7, 8]])
new_column = [3, 6, 9]
df.append_column("c", new_column)
self.assertEqual(df.get_column("c"), [3, 6, 9])
# test if new column has the wrong length
new_column = [3, 6, 9, 11]
with self.assertRaises(Exception) as cm:
df.append_column("d", new_column)
e = cm.exception
self.assertEqual(e.__str__(), "The new column has the wrong length.")