box-o-sand/leetcode/test_stuff.py

import pytest

import stuff
import stdlib


@pytest.mark.parametrize(
    ("n", "expected"),
    [
        (0, 0),
        (1, 1),
        (5, 2),
        (4, 2),
        (8, 2),
        (9, 3),
    ],
)
def test_find_sqrt_ish(n: int, expected: int):
    assert stuff.find_sqrt_ish(n) == expected


@pytest.mark.parametrize(
    ("n", "expected"),
    [
        (3, "III"),
        (58, "LVIII"),
        (1994, "MCMXCIV"),
    ],
)
def test_int_to_roman(n: int, expected: str):
    assert stuff.Roman.i2r(n) == expected


@pytest.mark.parametrize(
    ("ops", "expected"),
    [
        (
            (
                [
                    ("new",),
                    ("push", -2),
                    ("push", 0),
                    ("push", -3),
                    ("getMin",),
                    ("pop",),
                    ("top",),
                    ("getMin",),
                ]
            ),
            [
                None,
                None,
                None,
                None,
                -3,
                None,
                0,
                -2,
            ],
        )
    ],
)
def test_min_stack(ops: list[tuple[str] | tuple[str, int]], expected: list[int | None]):
    returned: list[int | None] = []
    inst: stuff.MinStack | None = None

    for op in ops:
        if len(op) == 1:
            if op[0] == "new":
                inst = stuff.MinStack()
                returned.append(None)
                continue

            returned.append(getattr(inst, op[0])())
            continue

        method, arg = op
        returned.append(getattr(inst, method)(arg))

    assert returned == expected


@pytest.mark.parametrize(
    ("head", "expected"),
    [
        (None, None),
        (
            stdlib.ListNode(
                4, stdlib.ListNode(2, stdlib.ListNode(1, stdlib.ListNode(3)))
            ),
            stdlib.ListNode(
                1, stdlib.ListNode(2, stdlib.ListNode(3, stdlib.ListNode(4)))
            ),
        ),
        (
            stdlib.ListNode(
                4,
                stdlib.ListNode(
                    19,
                    stdlib.ListNode(
                        14,
                        stdlib.ListNode(
                            5,
                            stdlib.ListNode(
                                -3,
                                stdlib.ListNode(
                                    1,
                                    stdlib.ListNode(
                                        8,
                                        stdlib.ListNode(
                                            5, stdlib.ListNode(11, stdlib.ListNode(15))
                                        ),
                                    ),
                                ),
                            ),
                        ),
                    ),
                ),
            ),
            stdlib.ListNode(
                -3,
                stdlib.ListNode(
                    1,
                    stdlib.ListNode(
                        4,
                        stdlib.ListNode(
                            5,
                            stdlib.ListNode(
                                5,
                                stdlib.ListNode(
                                    8,
                                    stdlib.ListNode(
                                        11,
                                        stdlib.ListNode(
                                            14, stdlib.ListNode(15, stdlib.ListNode(19))
                                        ),
                                    ),
                                ),
                            ),
                        ),
                    ),
                ),
            ),
        ),
    ],
)
def test_sort_linked_list(head: stdlib.ListNode | None, expected: stdlib.ListNode | None):
    if head is None:
        assert stuff.sort_linked_list(head) == expected
        return

    assert stuff.linked_list_to_list(
        stuff.sort_linked_list(head)
    ) == stuff.linked_list_to_list(expected)


@pytest.mark.parametrize(
    ("root", "expected"),
    [
        (
            stdlib.Node(
                1,
                left=stdlib.Node(2, left=stdlib.Node(4), right=stdlib.Node(5)),
                right=stdlib.Node(3, right=stdlib.Node(7)),
            ),
            [1, None, 2, 3, None, 4, 5, 7, None],
        ),
    ],
)
def test_connect_binary_tree_right(
    root: stdlib.Node | None, expected: list[int | None] | None
):
    if expected is None:
        assert root is None
        return

    connected, serialized = stuff.connect_binary_tree_right(root)
    assert connected is not None
    assert serialized == expected