russell@unturf.com
8fe389d45e
Disk-aware strategy picker for the corpus reshard. Pure-compute API:
feed it (corpus facts, free bytes) and it returns a HydrationPlan
naming one of three strategies + a peak-draw estimate.
Strategies (preference order):
all_at_once - direct source→target with overlap; fastest;
peak draw = corpus + WAL × M_target + VACUUM
per_source_shard - pack→delete→hydrate per round; safest;
peak draw = pack + WAL + 1 target VACUUM
streaming_row - in-place mutation, --allow-in-place opt-in;
peak draw = WAL only
Safety budget: 4 GB margin above the strategy's peak draw — the
buffer that survives one bad sort + one badly-sized WAL grow.
Production-host preview at 95 GB free / 38 GB corpus:
strategy: all_at_once
peak draw est: 64.2 GB
free at peak: 31.2 GB (well above safety)
rationale: all_at_once: peak draw 64.2 GB + 4.3 GB safety ≤ free 95.4 GB
The plan's full readout (strategy, free bytes, peak draw, rationale)
is JSON-serializable so it goes into the migration audit event as a
single forensic record.
14 tests cover: strategy pick by free-disk level, --force_strategy
override, --allow-in-place gating, skewed shard sizes, json-
serializable audit body, target_M validation.
No execution yet — this is just the planner. Next: the strategy A
executor (direct source→target read+write).
223 lines
6.9 KiB
Python
223 lines
6.9 KiB
Python
"""Tests for the #000065 hydration planner.
|
||
|
||
Pure-compute tests — never touches SQLite or the network. Validates
|
||
that the planner picks the right strategy from a given (corpus size,
|
||
free disk) pair, and that the chosen strategy's peak-draw estimate
|
||
matches the safety budget.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
from pathlib import Path
|
||
|
||
import pytest
|
||
|
||
from arborist.migrate import (
|
||
CorpusFacts,
|
||
HydrationPlan,
|
||
HydrationPlanner,
|
||
InsufficientDisk,
|
||
SAFETY_MARGIN_BYTES,
|
||
)
|
||
|
||
|
||
GB = 1024 ** 3
|
||
|
||
|
||
@pytest.fixture
|
||
def four_shard_corpus(tmp_path: Path) -> Path:
|
||
"""Create 4 empty 'shard' files of realistic sizes (~8.8 GB each)."""
|
||
shards_dir = tmp_path / "shards"
|
||
shards_dir.mkdir()
|
||
# Don't allocate real bytes; touch + truncate is enough for the
|
||
# planner since it reads ``.stat().st_size``.
|
||
sizes_gb = [8.8, 8.8, 8.8, 8.8]
|
||
for i, size_gb in enumerate(sizes_gb):
|
||
path = shards_dir / f"{i:03d}.db"
|
||
path.touch()
|
||
os = __import__("os")
|
||
os.truncate(path, int(size_gb * GB))
|
||
return shards_dir
|
||
|
||
|
||
@pytest.fixture
|
||
def target_dir(tmp_path: Path) -> Path:
|
||
d = tmp_path / "shards.v2"
|
||
d.mkdir()
|
||
return d
|
||
|
||
|
||
class TestCorpusFacts:
|
||
def test_reads_sizes(self, four_shard_corpus):
|
||
f = CorpusFacts.from_dir(four_shard_corpus)
|
||
assert len(f.shard_files) == 4
|
||
assert f.total_bytes == int(8.8 * GB) * 4
|
||
assert f.largest_shard_bytes == int(8.8 * GB)
|
||
|
||
def test_empty_dir_rejected(self, tmp_path):
|
||
empty = tmp_path / "empty"
|
||
empty.mkdir()
|
||
with pytest.raises(ValueError, match="no shard files"):
|
||
CorpusFacts.from_dir(empty)
|
||
|
||
|
||
class TestStrategyPick:
|
||
def test_88_gb_free_picks_all_at_once(
|
||
self, four_shard_corpus, target_dir
|
||
):
|
||
"""The production-host case: 88 GB free, 35 GB corpus."""
|
||
plan = HydrationPlanner().plan(
|
||
source_dir=four_shard_corpus,
|
||
target_dir=target_dir,
|
||
target_M=4,
|
||
free_bytes=88 * GB,
|
||
)
|
||
assert plan.strategy == "all_at_once"
|
||
|
||
def test_47_gb_free_falls_to_per_shard(
|
||
self, four_shard_corpus, target_dir
|
||
):
|
||
"""The old-host case: 47 GB free, 35 GB corpus."""
|
||
plan = HydrationPlanner().plan(
|
||
source_dir=four_shard_corpus,
|
||
target_dir=target_dir,
|
||
target_M=4,
|
||
free_bytes=47 * GB,
|
||
)
|
||
assert plan.strategy == "per_source_shard"
|
||
|
||
def test_15_gb_free_refused_without_in_place(
|
||
self, four_shard_corpus, target_dir
|
||
):
|
||
"""At 15 GB free, no opt-in: strategy B doesn't fit either."""
|
||
with pytest.raises(InsufficientDisk):
|
||
HydrationPlanner().plan(
|
||
source_dir=four_shard_corpus,
|
||
target_dir=target_dir,
|
||
target_M=4,
|
||
free_bytes=15 * GB,
|
||
)
|
||
|
||
def test_10_gb_free_with_in_place_picks_streaming(
|
||
self, four_shard_corpus, target_dir
|
||
):
|
||
"""At 10 GB free with --allow-in-place: strategy C (streaming)."""
|
||
plan = HydrationPlanner().plan(
|
||
source_dir=four_shard_corpus,
|
||
target_dir=target_dir,
|
||
target_M=4,
|
||
free_bytes=10 * GB,
|
||
allow_in_place=True,
|
||
)
|
||
assert plan.strategy == "streaming_row"
|
||
|
||
def test_zero_free_refused_even_with_in_place(
|
||
self, four_shard_corpus, target_dir
|
||
):
|
||
with pytest.raises(InsufficientDisk):
|
||
HydrationPlanner().plan(
|
||
source_dir=four_shard_corpus,
|
||
target_dir=target_dir,
|
||
target_M=4,
|
||
free_bytes=1 * GB,
|
||
allow_in_place=True,
|
||
)
|
||
|
||
|
||
class TestForcedStrategy:
|
||
def test_force_strategy_b_even_when_a_fits(
|
||
self, four_shard_corpus, target_dir
|
||
):
|
||
"""Operator override: pick B even if A fits."""
|
||
plan = HydrationPlanner().plan(
|
||
source_dir=four_shard_corpus,
|
||
target_dir=target_dir,
|
||
target_M=4,
|
||
free_bytes=88 * GB,
|
||
force_strategy="per_source_shard",
|
||
)
|
||
assert plan.strategy == "per_source_shard"
|
||
assert "forced" in plan.rationale
|
||
|
||
def test_force_unavailable_strategy_rejected(
|
||
self, four_shard_corpus, target_dir
|
||
):
|
||
with pytest.raises(ValueError, match="not available"):
|
||
HydrationPlanner().plan(
|
||
source_dir=four_shard_corpus,
|
||
target_dir=target_dir,
|
||
target_M=4,
|
||
free_bytes=88 * GB,
|
||
force_strategy="streaming_row", # not allow_in_place
|
||
)
|
||
|
||
|
||
class TestPlanFields:
|
||
def test_audit_body_is_json_serializable(
|
||
self, four_shard_corpus, target_dir
|
||
):
|
||
import json
|
||
plan = HydrationPlanner().plan(
|
||
source_dir=four_shard_corpus,
|
||
target_dir=target_dir,
|
||
target_M=4,
|
||
free_bytes=88 * GB,
|
||
)
|
||
json.dumps(plan.as_audit_body()) # no exception
|
||
|
||
def test_free_at_peak_reasonable(
|
||
self, four_shard_corpus, target_dir
|
||
):
|
||
plan = HydrationPlanner().plan(
|
||
source_dir=four_shard_corpus,
|
||
target_dir=target_dir,
|
||
target_M=4,
|
||
free_bytes=88 * GB,
|
||
)
|
||
# At least the safety margin remains free at peak
|
||
# (otherwise the planner would've rejected this strategy).
|
||
assert plan.free_at_peak_bytes() >= SAFETY_MARGIN_BYTES
|
||
|
||
|
||
class TestTargetMValidation:
|
||
def test_target_M_zero_rejected(self, four_shard_corpus, target_dir):
|
||
with pytest.raises(ValueError, match="target_M"):
|
||
HydrationPlanner().plan(
|
||
source_dir=four_shard_corpus,
|
||
target_dir=target_dir,
|
||
target_M=0,
|
||
free_bytes=88 * GB,
|
||
)
|
||
|
||
def test_target_M_one(self, four_shard_corpus, target_dir):
|
||
"""M=1 (consolidate to a single shard) is a valid target."""
|
||
plan = HydrationPlanner().plan(
|
||
source_dir=four_shard_corpus,
|
||
target_dir=target_dir,
|
||
target_M=1,
|
||
free_bytes=88 * GB,
|
||
)
|
||
assert plan.target_M == 1
|
||
|
||
|
||
class TestVariedShardSizes:
|
||
"""Skew test: one shard much larger than the others."""
|
||
|
||
def test_strategy_b_bounded_by_largest(self, tmp_path):
|
||
import os
|
||
shards = tmp_path / "shards"
|
||
shards.mkdir()
|
||
# 3 small + 1 large.
|
||
sizes = [1 * GB, 1 * GB, 1 * GB, 20 * GB]
|
||
for i, s in enumerate(sizes):
|
||
p = shards / f"{i:03d}.db"
|
||
p.touch()
|
||
os.truncate(p, s)
|
||
plan = HydrationPlanner().plan(
|
||
source_dir=shards,
|
||
target_dir=tmp_path / "out",
|
||
target_M=4,
|
||
free_bytes=40 * GB,
|
||
)
|
||
# Strategy A peak should fail (23 GB new + 4 × 4 WAL + scratch >40),
|
||
# so we expect strategy B which is bounded by the 20 GB shard.
|
||
assert plan.strategy == "per_source_shard"
|