python/test_properties.py

from avl import AVLTree
from wavl import WAVLTree
from ravl import RAVLTree
from rbt import RBTree

import logging
import random

import hypothesis
import hypothesis.strategies as st
import pytest

logger = logging.getLogger(__name__)


@pytest.mark.parametrize(
    ("RankBalancedTree"), [RBTree, AVLTree, WAVLTree, RAVLTree]
)
@hypothesis.settings(max_examples=10000)
@hypothesis.given(values=st.sets(st.integers()))
def test_insert_random(RankBalancedTree, values):
    tree = RankBalancedTree()

    for value in values:
        tree.insert(value)
        assert tree.is_correct
        assert tree.search(value) is not None

    assert all(tree.search(key) is not None for key in values)


@st.composite
def delete_strategy(draw):
    values = list(draw(st.sets(st.integers())))
    delete_order = values.copy()
    random.shuffle(delete_order)
    return (values, delete_order)


def _report(t_before: str, t_after: str) -> None:
    with open("before.dot", "w") as before, open("after.dot", "w") as after:
        print(t_before, file=before)
        print(t_after, file=after)


@pytest.mark.parametrize(
    ("RankBalancedTree"), [RBTree, AVLTree, WAVLTree, RAVLTree]
)
@hypothesis.settings(max_examples=10000)
@hypothesis.given(config=delete_strategy())
def test_delete_random(RankBalancedTree, config):
    values, delete_order = config

    tree = RankBalancedTree()

    for value in values:
        tree.insert(value)

    for value in delete_order:
        before = str(tree)
        tree.delete(value)
        after = str(tree)

        try:
            assert tree.is_correct
            assert tree.search(value) is None
        except AssertionError:
            logger.info(
                f"[FAIL] Delete {value} from {values} in order {delete_order}"
            )
            _report(before, after)
            logger.info(f"Before:\n{before}")
            logger.info(f"After:\n{after}")
            raise