python/test_avl.py

from avl import AVLTree, _balance_factor

import logging
import random
from typing import List

import hypothesis
import hypothesis.strategies as st
import pytest

logger = logging.getLogger(__name__)


@pytest.mark.parametrize(
    ("values", "expected_balance_factor"),
    [
        ([], 0),
        ([1], 0),
        ([1, -1], -1),
        ([1, 2], 1),
    ],
)
def test_balance_factor(
    values: List[int], expected_balance_factor: int
) -> None:
    tree = AVLTree()
    for v in values:
        tree.insert(v)

    assert _balance_factor(tree.root) == expected_balance_factor


def test_incorrect_avl() -> None:
    tree = AVLTree()
    tree.insert(0)
    tree.root.rank = 1
    assert not tree.is_correct


def test_empty() -> None:
    tree = AVLTree()

    assert tree.root is None
    assert tree.is_correct


def test_one_node() -> None:
    tree = AVLTree()
    tree.insert(1)

    assert tree.root is not None
    assert 1 == tree.root.value
    assert 0 == tree.root.rank
    assert tree.root.left is None
    assert tree.root.right is None

    assert tree.is_correct


@pytest.mark.parametrize("values", [[1, 2], [1, 2, 0]])
def test_no_rebalance_needed(values):
    tree = AVLTree()

    for value in values:
        tree.insert(value)
        assert tree.is_correct


def test_three_nodes_rebalanced():
    tree = AVLTree()

    for value in (1, 2, 3):
        print(tree)
        tree.insert(value)

    assert tree.is_correct


def test_bigger_tree():
    tree = AVLTree()

    for i in range(50):
        tree.insert(i)
        assert tree.is_correct


def test_bigger_tree_reversed():
    tree = AVLTree()

    for i in range(50):
        tree.insert(-i)
        assert tree.is_correct


def test_promote():
    tree = AVLTree()

    for value in (0, 1, -1):
        tree.insert(value)
        assert tree.is_correct


@pytest.mark.parametrize(
    "values",
    [
        [0, 1, -2, -1, -3],
        [0, 1, -2, -1, -3, -4, 4, 9, 7],
        [0, 1, -2, -1, -3, -4, 4, 9, 7, 5, -5, 8],
    ],
)
def test_rotate(values):
    tree = AVLTree()

    for value in values:
        tree.insert(value)
        assert tree.is_correct


@hypothesis.settings(max_examples=1000, deadline=None)
@hypothesis.given(values=st.sets(st.integers()))
def test_insert_random(values):
    tree = AVLTree()

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


@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)


@hypothesis.settings(max_examples=1000, deadline=None)
@hypothesis.given(config=delete_strategy())
def test_delete_random(config):
    values, delete_order = config

    tree = AVLTree()

    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
        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


# @unittest.skip("Only for replicating hypothesis")
@pytest.mark.parametrize(
    "values, delete_order",
    [
        ([4, 2, 5, 0, 3], [5, 0, 4, 2, 3]),
        ([1, 2, 3, 4, 5], [5]),
        ([0, 1, 2, -15, 17, -2, -1], [17, 2, 1, 0, -2, -15, -1]),
        (
            [
                0,
                1,
                2,
                3,
                4,
                128,
                1536,
                7,
                25223,
                1159,
                16,
                32915,
                148,
                147,
                24,
                26,
                220,
                -1,
                -9,
                -8,
                -6,
                -5,
                -4,
                -3,
                -2,
            ],
            [
                24,
                0,
                -3,
                -9,
                1536,
                128,
                7,
                16,
                -4,
                -6,
                4,
                2,
                -5,
                3,
                -1,
                32915,
                -8,
                1159,
                -2,
                26,
                220,
                25223,
                1,
                147,
                148,
            ],
        ),
    ],
)
def test_delete_minimal(values, delete_order):
    tree = AVLTree()

    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
        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