def test_print_tree(self): tree = Tree(None, None) tree._gen = lambda x: [(0, 12345, 'ref1'), (1, 23456, 'ref2')] with patch('sys.stdout', StringIO()) as stdout: tree.print_tree() self.assertEqual(stdout.getvalue(), " 12345 ref1\n" " 23456 ref2\n")
def test_walk_max_results(self): tree = Tree(None, None) tree._gen = lambda x: list(range(20)) res = list(tree.walk(maxresults=3)) self.assertEqual(res, [0, 1, 2, (0, 0, "==== Max results reached ====")])