Modified the test_tree.py