arborist/tests/test_store_migration_memoization.py

"""Regression tests for ticket #000026 Phase 1.

Per-process migration memoization. ``arborist.store.connect()`` used
to run every forward-migration probe on every open — hundreds of
redundant probe sequences per typical query. With memoization, the
migration block runs once per (physical file, process) and is skipped
on every subsequent open of the same shard.

Tests pin:

- First connect runs migrations.
- Second connect on the same path skips them.
- Functional integrity is preserved (all schema is in place after
  the second connect).
- The cache invalidates when a file is replaced at the same path
  (different inode → migrations re-run).
- ``_clear_migration_cache()`` resets the memo for tests / explicit
  re-probe.
"""

from __future__ import annotations

from pathlib import Path

from arborist import store
from arborist.store import (
    _clear_migration_cache,
    connect,
    invalidate_migration_cache,
)


class _MigrationProbeCounter:
    """Wraps each migration helper and counts invocations.

    Patches the module-level names so ``connect()`` (which references
    them via the ``arborist.store`` module globals) calls the
    counted versions.
    """

    # Mirrors the forward-migration sequence in ``store.connect()`` — keep
    # in sync when a migration is added (otherwise its memoization goes
    # unverified, and the count assertions below silently under-count).
    NAMES = (
        "_migrate_audit_mode",
        "_migrate_mesh_peer_chains",
        "_migrate_document_http_meta",
        "_migrate_selfmodel_tables",
        "_migrate_capital_ledger",
        "_migrate_memory_root",
        "_migrate_adapter_loss_reports",
        "_migrate_controller_events",
        "_migrate_fork_score_branches",
    )

    def __init__(self, monkeypatch):
        self.counts: dict[str, int] = {n: 0 for n in self.NAMES}
        for name in self.NAMES:
            original = getattr(store, name)

            def _wrapped(conn, *, _orig=original, _name=name):
                self.counts[_name] += 1
                return _orig(conn)

            monkeypatch.setattr(store, name, _wrapped)

    @property
    def total(self) -> int:
        return sum(self.counts.values())


def test_first_connect_runs_all_forward_migrations(tmp_path: Path, monkeypatch):
    _clear_migration_cache()
    counter = _MigrationProbeCounter(monkeypatch)
    db = tmp_path / "first.db"

    conn = connect(db)
    conn.close()

    # Every probe ran exactly once.
    assert counter.total == len(counter.NAMES)
    for name in counter.NAMES:
        assert counter.counts[name] == 1


def test_second_connect_skips_migrations(tmp_path: Path, monkeypatch):
    _clear_migration_cache()
    db = tmp_path / "second.db"

    # Warm the cache with a fresh open (no probe counter yet).
    conn = connect(db)
    conn.close()

    # Now patch in counters and open again. Zero probes should run.
    counter = _MigrationProbeCounter(monkeypatch)
    conn = connect(db)
    conn.close()
    assert counter.total == 0


def test_second_connect_schema_is_intact(tmp_path: Path):
    """Belt-and-suspenders: migration skipping must not break schema
    visibility. Open + close + reopen, then read every table the
    migrations create."""
    _clear_migration_cache()
    db = tmp_path / "intact.db"
    connect(db).close()

    conn = connect(db)
    try:
        # Smoke-check tables that each migration is responsible for.
        for table in (
            "providence_cache",       # _migrate_audit_mode
            "mesh_peer_chains",       # _migrate_mesh_peer_chains
            "document_http_meta",     # _migrate_document_http_meta
            "selfmodel_records",      # _migrate_selfmodel_tables
            "capital_ledger",         # _migrate_capital_ledger
            "memory_records",         # _migrate_memory_root
            "adapter_loss_reports",   # _migrate_adapter_loss_reports
            "controller_events",      # _migrate_controller_events
            "fork_score_branches",    # _migrate_fork_score_branches
        ):
            row = conn.execute(
                "SELECT name FROM sqlite_master WHERE type='table' AND name=?",
                (table,),
            ).fetchone()
            assert row is not None, f"missing table after re-open: {table}"
    finally:
        conn.close()


def test_clear_migration_cache_forces_reprobe(tmp_path: Path, monkeypatch):
    _clear_migration_cache()
    db = tmp_path / "clear.db"

    connect(db).close()

    counter = _MigrationProbeCounter(monkeypatch)

    # Without clearing — zero probes.
    connect(db).close()
    assert counter.total == 0

    # After clearing — full re-probe.
    _clear_migration_cache()
    connect(db).close()
    assert counter.total == len(counter.NAMES)


def test_invalidate_migration_cache_after_file_replacement(
    tmp_path: Path, monkeypatch
):
    """Caller-driven invalidation contract: after replacing a shard
    at the same path, the caller must call
    ``invalidate_migration_cache(path)`` so the next ``connect()``
    re-runs the full probe sequence on the new file.

    We don't auto-detect replacement — (dev, inode) is unreliable
    under tmpfs inode reuse, and (mtime, size) drifts naturally as
    SQLite operates. Path-only with explicit invalidation is the
    honest contract; tests + snapshot-restore flows are responsible
    for calling the invalidator."""
    _clear_migration_cache()
    db = tmp_path / "replaced.db"

    # First open — populate cache.
    connect(db).close()

    # Replace the file at the same path.
    db.unlink()

    # Without invalidation, the cache still claims migrations are
    # current — schema would be silently absent on the new file.
    invalidate_migration_cache(db)

    counter = _MigrationProbeCounter(monkeypatch)
    conn = connect(db)
    conn.close()
    assert counter.total == len(counter.NAMES)


def test_distinct_paths_each_get_one_probe(tmp_path: Path, monkeypatch):
    """Two different shard files get one probe each — the cache is
    per-path, not global."""
    _clear_migration_cache()
    db_a = tmp_path / "a.db"
    db_b = tmp_path / "b.db"

    counter = _MigrationProbeCounter(monkeypatch)
    connect(db_a).close()
    connect(db_b).close()
    # Two paths × the full migration sequence.
    assert counter.total == 2 * len(counter.NAMES)

    # Re-open both — cache hit on each, zero new probes.
    before = counter.total
    connect(db_a).close()
    connect(db_b).close()
    assert counter.total == before