Red-Black tree fuzzy tester

rbtest.py

#!/usr/bin/env python3

TEST_ELEM_COUNT  = 20
TEST_FAILED_FILE = 'rbtest_failed.txt'

import random, sys, traceback, ast


import rb # Pas aan indien nodig

def new_tree():
    # Pas aan
    return rb.RBT()

def get_key(n):
    # Pas aan
    return n.key

def get_left(n):
    # Pas aan
    return n.left

def get_right(n):
    # Pas aan
    return n.right

def is_black(n):
    # Pas aan
    return rb.is_black(n)

def is_red(n):
    # Pas aan
    return rb.is_red(n)

def is_nil(n):
    # Pas aan
    return n is None

def get_root(t):
    # Pas aan
    return t.root

def insert(t, k):
    # Pas aan
    t.searchTreeInsert(rb.Node(k, None))

def remove(t, k):
    # Pas aan
    t.searchTreeDelete(k)


def print_tree(t):
    def print_node(n, lvl = 0):
        if n is None:
            return
        if get_left(n) is not None:
            print_node(get_left(n), lvl + 1)
        print('  ' * lvl, 'R:' if is_red(n) else 'B:', get_key(n), sep='')
        if get_right(n) is not None:
            print_node(get_right(n), lvl + 1)
    print_node(get_root(t))

def collect_tree(t):
    l = []
    t.inorderTraverse(lambda n: l.append(get_key(n)))
    return l

def dump_tree(msg, t, l):
    print(msg)
    print('Tree:')
    print_tree(t)
    print('Expected:', sorted(l))
    print('Got:     ', collect_tree(t))
    print()
    return False

def check_valid(t, l):
    ct = collect_tree(t)

    if set(l) != set(ct):
        return dump_tree('Elements are missing!', t, l)
    if sorted(l) != ct:
        return dump_tree('Elements are not in order!', t, l)

    if not is_black(get_root(t)):
        return dump_tree('Root is not black!', t, l)

    def iterate(node, black_count, black_counted, parent_is_red):
        if is_nil(node):
            if black_counted[0] is None:
                black_counted[0] = black_count
            elif black_counted[0] != black_count:
                return 'Tree is unbalanced!'
            return None if is_black(node) else 'NIL leaf is not black!'

        if is_red(node) and parent_is_red:
            return 'Red node has red parent!'

        if is_black(node):
            black_count += 1

        ret = iterate(get_left(node), black_count, black_counted, is_red(node))
        if ret is not None:
            return ret
        return iterate(get_right(node), black_count, black_counted, is_red(node))

    ret = iterate(get_root(t), 0, [None], False)
    if ret is not None:
        return dump_tree(ret, t, l)

    return True

def test(li, lr, check_before = None):
    state = 0
    try:
        t = new_tree()

        for i, e in enumerate(li):
            state += 1
            if check_before == state:
                dump_tree('State before error:', t, li[:i])
                print('Will insert', e)
            insert(t, e)
            if not check_valid(t, li[:i + 1]):
                print('Error during insertion of', e)
                return state

        lr = lr[:]
        while len(lr) > 0:
            state += 1
            if check_before == state:
                dump_tree('State before error:', t, lr)
                print('Will remove', lr[-1])
            e = lr.pop()
            remove(t, e)
            if not check_valid(t, lr):
                print('Error during removal of', e)
                return state

        return None

    except Exception as e:
        print('Exception during testing!')
        print('Tree:')
        print_tree(t)
        traceback.print_exc()
        return state

def save_test(li, lr, result):
    print('Writing test case to', TEST_FAILED_FILE)
    with open(TEST_FAILED_FILE, 'w') as f:
        f.write(str((li, lr, result)))

if __name__ == '__main__':
    try:
        with open(TEST_FAILED_FILE) as f:
            print('Testing case in', TEST_FAILED_FILE)
            li, lr, result = ast.literal_eval(f.read())
            result = test(li, lr, result)
            if result is None:
                print('Test passed! You can remove', TEST_FAILED_FILE)
                sys.exit(0)
            else:
                print('Test failed!')
                save_test(li, lr, result)
                sys.exit(1)
    except IOError:
        pass

    print('Running fuzzy tests...')
    n = 0
    while True:
        li = [*range(TEST_ELEM_COUNT)] # insert list
        lr = [*range(TEST_ELEM_COUNT)] # remove list
        random.shuffle(li)
        random.shuffle(lr)
        result = test(li, lr)
        if result is not None:
            print('Test', n , 'failed!')
            save_test(li, lr, result)
            sys.exit(1)
        n += 1
        if n % 1000 == 0:
            print('Tested', n, 'cases')