"""Stdlib-only tests (no pytest): plain asserts in functions called from a __main__ block. Run: `python3 test_btree.py`. The headline test is test_sorted_load_builds_a_worse_index — the aha from commentary.html section 6, shrunk to 20,000 keys so it runs in a couple of seconds. The rest pin the page format, the split point, and the promise that a lookup costs exactly `height` page reads. Index files go in a temp directory and are cleaned up. """ import hashlib import math import os import random import shutil import tempfile from btree import MAX_KEYS, PAGE_SIZE, BPlusTree, Node, Pager, scan TMP = tempfile.mkdtemp(prefix="btree-test-") def build(keys, name): """Build an index over `keys` in insertion order. Returns (pager, tree).""" pager = Pager(os.path.join(TMP, name + ".idx")) tree = BPlusTree(pager) for key in keys: tree.insert(key, key * 10) return pager, tree def shuffled(n, seed=42): keys = list(range(n)) random.Random(seed).shuffle(keys) return keys def test_a_node_serializes_to_exactly_one_page(): """Every page is PAGE_SIZE bytes, and unpack(pack(x)) == x for both kinds.""" leaf = Node(True, [3, 9, 27], [30, 90, 270]) inner = Node(False, [10, 20], [7, 8, 9]) for node in (leaf, inner): buf = node.pack() assert len(buf) == PAGE_SIZE, len(buf) back = Node.unpack(buf) assert back.leaf == node.leaf assert back.keys == node.keys assert back.kids == node.kids # A full node still fits: 7 bytes of header + 63 slots of 16 bytes = 1015. full = Node(True, list(range(MAX_KEYS)), list(range(MAX_KEYS))) assert len(full.pack()) == PAGE_SIZE def test_scan_returns_the_slot_and_counts_its_comparisons(): """Leaves scan for where the key belongs; internal nodes for which child.""" keys = [10, 20, 30] assert scan(keys, 20, True) == (1, 2) # leaf: first slot with key >= 20 assert scan(keys, 20, False) == (2, 3) # inside: 20 lives in child 2 assert scan(keys, 5, True) == (0, 1) # stops on the first comparison assert scan(keys, 99, True) == (3, 3) # ran off the end: 3 comparisons def test_every_key_is_retrievable_and_absent_keys_are_not(): """5,000 keys inserted in scrambled order; all 5,000 come back.""" _, tree = build(shuffled(5_000), "roundtrip") assert all(tree.get(k) == k * 10 for k in range(5_000)) assert tree.get(-1) is None assert tree.get(5_000) is None assert tree.get(2_500_000) is None def test_the_64th_key_splits_the_root_exactly_in_half(): """MAX_KEYS=63, so the 64th key overflows the root leaf and halves it.""" pager, tree = build(range(63), "sixty_three") assert (tree.height, pager.next_pid) == (1, 1) pager, tree = build(range(64), "sixty_four") assert (tree.height, pager.next_pid) == (2, 3) # new root + two leaves assert tree.growth == [(64, 2)] root = Node.unpack(pager.read_page(tree.root)) assert root.keys == [32] and len(root.kids) == 2 halves = [len(Node.unpack(pager.read_page(pid)).keys) for pid in root.kids] assert halves == [32, 32] # The separator was *copied* up, not moved: key 32 is still in a leaf. assert tree.get(32) == 320 def test_a_lookup_reads_exactly_height_pages(): """The whole promise of the structure, at three different heights.""" for n, expected_height in ((63, 1), (2_000, 2), (20_000, 3)): pager, tree = build(shuffled(n, seed=1), f"height{n}") assert tree.height == expected_height, (n, tree.height) pager.reads = 0 assert tree.get(n // 2) == (n // 2) * 10 assert pager.reads == expected_height, (n, pager.reads) assert len(tree.last_path) == expected_height def test_reinserting_a_key_updates_it_in_place(): """An update rewrites one page; it does not allocate or grow the tree.""" pager, tree = build(shuffled(1_000), "update") before_pages, before_height = pager.next_pid, tree.height pager.writes = 0 tree.insert(500, -1) assert tree.get(500) == -1 assert pager.next_pid == before_pages assert tree.height == before_height assert pager.writes == 1 def test_sorted_load_builds_a_worse_index(): """The aha, at 20,000 keys: ascending inserts leave every leaf half full. An ascending load only ever overflows the rightmost leaf, and a 50/50 split leaves the left half at 32/63 forever. Random inserts split leaves all over the tree and settle at ln 2 -- the classic B-tree occupancy. """ _, asc = build(range(20_000), "asc") _, rnd = build(shuffled(20_000), "rnd") a, r = asc.stats(), rnd.stats() assert a["fill"] == 32 / MAX_KEYS # exactly 50.79% assert abs(r["fill"] - math.log(2)) < 1e-4 # 69.31%, i.e. ln 2 assert a["leaves"] > r["leaves"] == 458 assert a["pages"] == 644 and r["pages"] == 467 # 38% more pages # Both are height 3 at this N: the extra *level* is a threshold effect # that needs ~68,600 ascending keys (commentary section 6.4). assert a["height"] == r["height"] == 3 assert asc.growth == [(64, 2), (2080, 3)] def test_the_index_file_is_byte_identical_across_builds(): """Page ids come from a monotonic allocator, so the bytes are reproducible.""" digests = [] for run in range(2): pager, _ = build(shuffled(5_000), f"repeat{run}") pager.f.flush() with open(pager.path, "rb") as f: digests.append(hashlib.sha256(f.read()).hexdigest()) assert digests[0] == digests[1], digests TESTS = [ test_a_node_serializes_to_exactly_one_page, test_scan_returns_the_slot_and_counts_its_comparisons, test_every_key_is_retrievable_and_absent_keys_are_not, test_the_64th_key_splits_the_root_exactly_in_half, test_a_lookup_reads_exactly_height_pages, test_reinserting_a_key_updates_it_in_place, test_sorted_load_builds_a_worse_index, test_the_index_file_is_byte_identical_across_builds, ] if __name__ == "__main__": failed = 0 try: for test in TESTS: try: test() except AssertionError as exc: failed += 1 print(f"FAIL {test.__name__}: {exc}") else: print(f"PASS {test.__name__}") finally: shutil.rmtree(TMP, ignore_errors=True) if failed: print(f"\n{failed} of {len(TESTS)} tests FAILED") raise SystemExit(1) print(f"\nAll {len(TESTS)} tests PASSED")