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.")