Python树的单元测试

TreeTest.py

class TreeNode:
    def __init__(self, data, left=None, right=None):
        self.data = data
        self.left = left
        self.right = right

def test_whole_group():
    r"""
    shape:
            10
         /      \
        6        14
       /\        /\
      4  8     12  16
    """
    node4 = TreeNode(4)
    node8 = TreeNode(8)
    node6 = TreeNode(6, node4, node8)
    node12 = TreeNode(12)
    node16 = TreeNode(16)
    node14 = TreeNode(14, node12, node16)
    node10 = TreeNode(10, node6, node14)
    return node10

def test_left_group():
    """
    shape:
                5
              /
             4
            /
           3
          /
         2
        /
       1
    """
    node1 = TreeNode(1)
    node2 = TreeNode(2, node1)
    node3 = TreeNode(3, node2)
    node4 = TreeNode(4, node3)
    node5 = TreeNode(5, node4)
    return node5

def test_right_group():
    r"""
    shape:
     1
      \
       2
        \
         3
          \
           4
            \
             5
    """
    node5 = TreeNode(5)
    node4 = TreeNode(4, right=node5)
    node3 = TreeNode(3, right=node4)
    node2 = TreeNode(2, right=node3)
    node1 = TreeNode(1, right=node2)
    return node1

def test_one():
    node = TreeNode(1)
    return node

def test_null():
    node = None
    return node

if __name__ == "__main__":
    whole = test_whole_group()
    left = test_left_group()
    right = test_right_group()
    one = test_one()
    null = test_null()