def test_build_tree(self):
leaves = [HuffmanNode(weight) for weight in range(1,5)]
root = HuffmanNode.build_tree(leaves)
# assertions based on calculating these by hand
self.assertEqual(root.weight, 10)
self.assertEqual(leaves[0].code, '110')
self.assertEqual(leaves[1].code, '111')
self.assertEqual(leaves[2].code, '10')
self.assertEqual(leaves[3].code, '0')
for leaf in leaves:
self.assertTrue(leaf.is_leaf)
def test_root_weight(self):
# sanity check that total weight should be sum of all node weights
for n in range(4, 100):
leaves = [HuffmanNode(weight) for weight in range(1, n)]
tree = HuffmanNode.build_tree(leaves)
self.assertEqual(tree.weight, sum(leaf.weight for leaf in leaves))