Shows how can we use unit tests while implementing heap related functions in Python

heap_example.py

def left(i):
    return 2 * i


def test_left():
    pass


def right(i):
    return 2 * i + 1


def test_right():
    pass


def parent(i):
    return i // 2


def test_parent():
    pass


def max_heapify(heap, heap_size, i):
    l = left(i)
    r = right(i)

    if l <= heap_size and heap[l] > heap[i]:
        largest = l
    else:
        largest = i
        
    if r <= heap_size and heap[r] > heap[largest]:
        largest = r

    if largest != i:
        heap[i], heap[largest] = heap[largest], heap[i]
        max_heapify(heap, heap_size, largest)


def test_max_heapify():
    test_cases = [
        {
            "name": "Simple Test Case",
            "input": [None, 1, 2, 3],
            "expected": [None, 3, 2, 1]
        },
        {
            "name": "Not so simple",
            "input": [None, 12, 3, 15, 1, 2, 17, 10],
            "expected": [None, 15, 3, 17, 1, 2, 12, 10]
        },
        {
            "name": "Empty List",
            "input": [None],
            "expected": [None]
        }
    ]
    
    for t in test_cases:
        max_heapify(t["input"], len(t["input"])-1, 1)
        assert t["expected"] == t["input"], t["name"]


def build_max_heap(heap):
    heap_size = len(heap) - 1
    for i in range(heap_size//2, 0, -1):
        max_heapify(heap, heap_size, i)


def test_build_max_heap():
    test_cases = [
        {
            "name": "test case 1",
            "input": [None, 12, 7, 1, 3, 10, 17, 19, 2, 5],
            "expected": [None, 19, 10, 17, 5, 7, 12, 1, 2, 3]
        }
    ]
    for t in test_cases:
        build_max_heap(t["input"])
        assert t["expected"] == t["input"], t["name"]


def heap_sort(heap):
    build_max_heap(heap)
    heap_size = len(heap) - 1
    for i in range(heap_size, 1, -1):
        heap[1], heap[i] = heap[i], heap[1]
        heap_size -= 1
        max_heapify(heap, heap_size, 1)


def tet_heap_sort():
    pass


def extract_max(heap, heap_size):
    max_item = heap[1]
    heap[1] = heap[heap_size]
    heap_size -= 1
    max_heapify(heap, heap_size, 1)
    return max_item


def test_extract_max():
    pass


if __name__ == "__main__":
    from binarytree import convert, pprint
    heap = [None, 12, 7, 1, 3, 10, 17, 19, 2, 5]
    print("Before building heap")
    pprint(convert(heap[1:]))
    build_max_heap(heap)
    print("After building heap")
    pprint(convert(heap[1:]))