def test_KeyList_Repeats(self):
tree = KeyedTree()
tree.add_pair("parent", "child")
tree.add_pair("parent", "brother")
tree.add_pair("parent", "sister")
tree.add_pair("grandparent", "parent")
expected = [
KeyedNode(n)
for n in ["parent", "child", "brother", "sister", "grandparent"]
]
nodes = []
tree.for_each_node(lambda n: nodes.append(n))
self.assertListEqual(nodes, expected)
class OrbitCounter:
def __init__(self):
self.tree = KeyedTree()
def add(self, orbiting, orbiter):
self.tree.add_pair(orbiting, orbiter)
def get_orbits(self):
count = 0
def counter(node):
nonlocal count
count += node.get_depth()
self.tree.for_each_node(counter)
return count
def get_hops(self, start_point, santa):
return len(self.tree.get_path(start_point, santa)) - 1
def test_ForEach_Empty(self):
nodes = []
tree = KeyedTree()
tree.for_each_node(lambda n: nodes.append(n))
self.assertListEqual(nodes, [])