arborist/tests/test_shard_routing.py
russell@unturf.com 3aae8119a6
#000065 step 1: routing helper + meta field + pre-migration snapshot
Three pieces, all read-only or additive — no shard mutation, no
schema-version bump:

1. shard_for_document(document_root, M) in arborist/document.py.
   Pure function: int(document_root[:8], 16) % M. 22 tests cover
   determinism, range-bounds, near-uniform distribution (±5pp at
   N=20k), and seven lock-in fixtures so peers will disagree
   loudly if anyone changes the formula.

2. corpus_shard_count meta field + get/set helpers in store.py.
   Lives in the existing key/value meta table; SCHEMA_VERSION
   stays at v9.8.0 (the DDL doesn't change and source_root is
   layout-independent, so cache records survive a reshard).
   Legacy shards (without the field) return None; reshard tool
   populates it on every target shard at migration time.

3. Pre-migration snapshot captured to
   bench/results/pre-migration-snapshot.json:
     docs        3,468,392  (3,468,226 globally unique)
     chunks      6,235,764
     edges      90,593,537
     audit       3,468,403
   This is the reference set post-reshard row counts must match.

4. Audit-event extraction script writes all 3.47M events from
   all 4 shards to /tmp/audit-events.ndjson (2.0 GB) for the
   Option-A canonical-chain consolidation step. Verifies chain
   integrity on extract — all 4 source chains report 0 breaks.

5. Fixed a wrong chunk count in docs/corpus-history.md
   (had ~3.54M/shard; actual is ~1.56M/shard) and added the
   edge-count column (~22.6M/shard, 90.6M total). 6.24M chunks
   total, not 14.12M.

Tests: 29 new pass (22 routing + 7 meta). No existing tests
touched.
2026-05-26 13:13:44 -04:00

101 lines
3.3 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

"""Tests for ``shard_for_document`` (#000065).
The routing helper picks a target shard index from a content hash.
Identical answer on every peer; uniform distribution across shards.
"""
from __future__ import annotations
import hashlib
import pytest
from arborist.document import shard_for_document
def _root(text: str) -> str:
"""Return a synthetic hex sha256 to use as a document_root."""
return hashlib.sha256(text.encode("utf-8")).hexdigest()
class TestDeterminism:
def test_same_input_same_output(self):
root = _root("the rain in spain")
assert shard_for_document(root, 4) == shard_for_document(root, 4)
def test_independent_of_prior_calls(self):
a = _root("a"), 4
b = _root("b"), 4
first = (shard_for_document(*a), shard_for_document(*b))
second = (shard_for_document(*a), shard_for_document(*b))
assert first == second
class TestBounds:
@pytest.mark.parametrize("M", [1, 2, 4, 8, 16, 64])
def test_within_range(self, M):
for i in range(200):
root = _root(f"doc-{i}")
idx = shard_for_document(root, M)
assert 0 <= idx < M
def test_M_equals_one_always_zero(self):
for i in range(50):
assert shard_for_document(_root(f"doc-{i}"), 1) == 0
def test_M_zero_rejected(self):
with pytest.raises(ValueError):
shard_for_document(_root("x"), 0)
def test_M_negative_rejected(self):
with pytest.raises(ValueError):
shard_for_document(_root("x"), -1)
class TestUniformity:
"""First-32-bit prefix of SHA-256 is uniform; check that distribution
is near-uniform across shards at meaningful sample sizes."""
@pytest.mark.parametrize("M", [2, 4, 8])
def test_distribution_within_5pct(self, M):
N = 20_000
counts = [0] * M
for i in range(N):
counts[shard_for_document(_root(f"doc-{i}"), M)] += 1
expected = N / M
for i, c in enumerate(counts):
ratio = c / expected
assert 0.95 <= ratio <= 1.05, (
f"shard {i}: {c} hits, expected ~{expected:.0f} "
f"({ratio:.3f}× — outside 5%)"
)
def test_M_changes_partition(self):
"""When M doubles, docs can move shard. Confirm at least one
sampled doc lands in different shards under different M values."""
roots = [_root(s) for s in ("YouTube", "SQLite", "arborist", "Wikipedia")]
any_varied = False
for r in roots:
seen = {shard_for_document(r, m) for m in (2, 4, 8, 16)}
if len(seen) >= 2:
any_varied = True
break
assert any_varied, "no sampled root produced different shards under varying M"
class TestKnownRoots:
"""Lock-in fixtures: if these change, peers will disagree."""
@pytest.mark.parametrize(
"root,M,expected",
[
("00000000" + "0" * 56, 4, 0),
("ffffffff" + "0" * 56, 4, 3),
("00000004" + "0" * 56, 4, 0),
("00000003" + "0" * 56, 4, 3),
("12345678" + "0" * 56, 4, 0),
("12345679" + "0" * 56, 4, 1),
("ffffffff" + "0" * 56, 1, 0),
],
)
def test_fixture(self, root, M, expected):
assert shard_for_document(root, M) == expected