def test_nodes_at_level(self): tree = BinaryDecisionTree(n_features=1) split = BinaryDecisionTreeSplit(feature_id=0, value=0.0) tree.split_node(0, split) tree.split_node(2, split) print(tree) self.assertEqual([0], tree.nodes_at_level(1)) self.assertEqual([1, 2], tree.nodes_at_level(2)) self.assertEqual([5, 6], tree.nodes_at_level(3)) self.assertEqual([], tree.nodes_at_level(4))