agnes-the-ai-analyst/tests/test_pull_invariants.py

"""Property-style invariant tests for ``cli.lib.pull_sync.run_stack_sync``.

Section 5.6 of the unified-stack design lists four hard invariants the
local store MUST satisfy after every pull, regardless of the prior state
or the sequence of mutations:

1. **No orphan parquets** — every file in ``<local>/data/_shared/`` has at
   least one live reference (symlink / hardlink / copy) pointing at it.
2. **No broken references** — every reference recorded in
   ``sync_state.json`` resolves to an actual ``_shared/<id>.parquet``.
3. **sync_state ↔ disk 1:1** — every state row has a disk file, every
   disk parquet (except orphans, see #1) has a state row.
4. **Memory bundles ↔ stack 1:1** — ``<local>/memory/<slug>/bundle.md``
   exists iff the user has the domain in their stack.

This test fires the real sync engine against a randomized but
deterministic (seeded RNG) sequence of manifest mutations: add packages,
drop packages, change md5s, swap tables between packages, add/remove
direct tables, add/remove memory domains. After every step the
``_assert_invariants`` helper walks the on-disk store + state and asserts
all four hold.

Manifest fetcher is a pure-Python stub: it writes deterministic parquet
bytes (table_id || md5 || row count) so md5 mismatches reliably trigger
re-fetch logic; the memory bundle fetcher returns a function of the
domain slug.
"""

from __future__ import annotations

import hashlib
import json
import os
import random
from pathlib import Path
from typing import Any, Dict, List, Tuple

import pytest

from cli.lib.pull_sync import (
    PullStackOptions,
    audit_invariants,
    run_stack_sync,
)


# ---------------------------------------------------------------------------
# Fake server-side state — deterministic parquet bytes per (id, md5)
# ---------------------------------------------------------------------------


def _fake_parquet_bytes(table_id: str, md5_hint: str) -> bytes:
    """Return a unique blob per (id, md5). Not real parquet — the sync
    layer only inspects file md5 + path; the actual decoder is unused in
    invariant tests."""
    return f"FAKE-PARQUET|id={table_id}|md5={md5_hint}".encode("utf-8")


def _make_fetcher(md5_index: Dict[str, str]):
    """Returns a (url, dest_path) fetcher closure. ``md5_index`` maps
    table_id → the md5 we want the freshly-written file to compute to.
    """

    def _fetch(url: str, dest: Path) -> None:
        # URL shape (legacy): ``/api/data/{id}/download``
        tid = url.rstrip("/").rsplit("/", 2)[-2]
        md5 = md5_index.get(tid, "")
        dest.parent.mkdir(parents=True, exist_ok=True)
        body = _fake_parquet_bytes(tid, md5)
        dest.write_bytes(body)

    return _fetch


def _md5_of(path: Path) -> str:
    return hashlib.md5(path.read_bytes()).hexdigest()


def _bundle_fetcher_factory(domain_bodies: Dict[str, bytes]):
    def _fetch(slug: str) -> bytes:
        return domain_bodies.get(slug, b"")

    return _fetch


# ---------------------------------------------------------------------------
# Manifest builders
# ---------------------------------------------------------------------------


def _make_table_entry(tid: str, name: str, md5: str) -> Dict[str, Any]:
    return {
        "id": tid,
        "name": name,
        "md5": md5,
        "query_mode": "local",
        "parquet_url": f"/api/data/{tid}/download",
    }


def _make_manifest(
    *,
    direct: List[Dict[str, Any]] | None = None,
    packages: List[Dict[str, Any]] | None = None,
    domains: List[Dict[str, Any]] | None = None,
) -> Dict[str, Any]:
    return {
        "direct_tables": direct or [],
        "data_packages": packages or [],
        "memory_domains": domains or [],
    }


# ---------------------------------------------------------------------------
# Invariant assertions
# ---------------------------------------------------------------------------


def _walk_refs(state: Dict[str, Any]):
    """Yield every (ref_path, shared_path, strategy) tuple recorded in
    sync_state. Handles both flat (direct_tables) and nested (packages)
    state layouts."""

    def _emit(entry):
        ref = entry.get("ref_path")
        shared = entry.get("shared_path")
        if ref:
            yield ref, shared, entry.get("strategy")

    for v in (state.get("direct_tables") or {}).values():
        if isinstance(v, dict):
            yield from _emit(v)

    for pkg in (state.get("data_packages") or {}).values():
        if isinstance(pkg, dict):
            for v in pkg.values():
                if isinstance(v, dict):
                    yield from _emit(v)


def _assert_invariants(local_dir: Path, expected_domains_in_stack: set[str]):
    """Walk disk + sync_state.json and assert all four invariants hold."""
    state_path = local_dir / "sync_state.json"
    state = json.loads(state_path.read_text()) if state_path.exists() else {}
    local_data = local_dir / "data"
    shared_dir = local_data / "_shared"

    # Collect every ref → shared mapping from state.
    refs_by_shared: dict[Path, list[Path]] = {}
    for ref, shared, _strategy in _walk_refs(state):
        ref_p = Path(ref)
        shared_p = Path(shared) if shared else None

        # Invariant 2: every state ref resolves to an existing _shared file.
        assert shared_p and shared_p.exists(), (
            f"broken reference: state ref {ref_p} → shared {shared_p} "
            f"missing on disk"
        )
        # The ref itself must also exist on disk (symlink/hardlink/copy).
        assert ref_p.exists() or ref_p.is_symlink(), (
            f"broken reference: ref {ref_p} not present on disk"
        )

        refs_by_shared.setdefault(shared_p.resolve(strict=False), []).append(ref_p)

    # Invariant 1: every file in _shared/ has at least one reference.
    if shared_dir.exists():
        for f in shared_dir.iterdir():
            if not f.is_file():
                continue
            refs = refs_by_shared.get(f.resolve(strict=False), [])
            assert refs, f"orphan _shared parquet: {f.name} has zero references"

    # Invariant 3: sync_state rows ↔ disk parquets are 1:1 inside data/.
    # Every ref recorded in state corresponds to a file present on disk
    # (covered above). The reverse direction: every parquet inside a
    # NON-``_shared`` data subdir must trace back to a state ref.
    disk_refs: list[Path] = []
    if local_data.exists():
        for sub in local_data.iterdir():
            if sub.name == "_shared" or not sub.is_dir():
                continue
            for f in sub.iterdir():
                if f.is_file() or f.is_symlink():
                    disk_refs.append(f.resolve(strict=False))
    state_refs = set()
    for ref, _, _ in _walk_refs(state):
        state_refs.add(Path(ref).resolve(strict=False))
    for f in disk_refs:
        # ``f`` resolves to the shared file via symlink; the on-disk
        # symlink path is what we compare against state.
        # Lookup by path string:
        ref_set = {Path(r).resolve(strict=False) for r in [ref for ref, _, _ in _walk_refs(state)]}
        # If the disk ref's PARENT.name is the dir + ref name matches a
        # state-recorded ref, it's accounted for. Use path equality on the
        # original (non-resolved) ref path saved in state, since on POSIX
        # resolving a symlink dereferences to the shared file.
        accounted = False
        for ref, _shared, _strat in _walk_refs(state):
            if Path(ref) == f or Path(ref).resolve(strict=False) == f:
                accounted = True
                break
            # On macOS the temp dir is symlinked (/private/var/...). Compare
            # absolute names with realpath on both sides.
            try:
                if os.path.realpath(ref) == os.path.realpath(str(f)):
                    accounted = True
                    break
            except OSError:
                pass
        assert accounted, f"untracked disk parquet under data/: {f}"

    # Invariant 4: memory bundle ↔ stack.
    memory_dir = local_dir / "memory"
    if memory_dir.exists():
        on_disk = {p.name for p in memory_dir.iterdir() if p.is_dir()}
    else:
        on_disk = set()
    assert on_disk == expected_domains_in_stack, (
        f"memory bundle/stack mismatch: disk={on_disk} stack={expected_domains_in_stack}"
    )

    # And the engine's own audit_invariants surface should be empty —
    # auto-heal already ran inside run_stack_sync but the audit walks the
    # post-write state.
    violations = audit_invariants(local_data, state)
    # Filter the "dangling shared" case for entries that the engine
    # cleaned up: audit_invariants is intentionally conservative and may
    # WARN on legitimate transient state during the test sequence. For
    # this property test we only fail on items NOT visible in the four
    # primary invariants above — `dangling shared` is allowed if the
    # corresponding state entry was already pruned.
    serious = [v for v in violations if not v.startswith("dangling shared")]
    assert not serious, f"audit_invariants flagged: {serious}"


# ---------------------------------------------------------------------------
# Scenario generator
# ---------------------------------------------------------------------------


def _gen_md5(tid: str, ver: int) -> str:
    return hashlib.md5(f"{tid}:v{ver}".encode()).hexdigest()


def _build_random_manifest(
    rng: random.Random,
    *,
    table_pool: List[str],
    version_clock: Dict[str, int],
) -> Dict[str, Any]:
    """Build a random manifest from a fixed pool of tables.

    Two-pass construction:

    1. **Decide which tables appear where, and finalize their version.**
       Bumping happens once per (table_id, manifest) so every appearance
       of the same id in this manifest carries the SAME md5. (Real
       servers maintain this invariant — a single table has one canonical
       md5 at any moment; a manifest that emits different md5s for the
       same id across direct/package entries is inconsistent and the
       sync engine can't reconcile it.)

    2. **Materialize the manifest entries.** All entries that reference
       the same tid share the same md5 from pass 1.
    """
    # ---- pass 1 — pick membership + finalize per-tid version ----
    direct_tids = rng.sample(table_pool, k=rng.randint(0, min(3, len(table_pool))))
    n_pkgs = rng.randint(0, 3)
    pkg_specs: List[Tuple[str, List[str]]] = []
    for i in range(n_pkgs):
        slug = f"pkg{i}"
        ptables = rng.sample(table_pool, k=rng.randint(1, min(3, len(table_pool))))
        pkg_specs.append((slug, ptables))

    referenced_tids = set(direct_tids)
    for _slug, ptables in pkg_specs:
        referenced_tids.update(ptables)

    final_md5: Dict[str, str] = {}
    for tid in sorted(referenced_tids):
        ver = version_clock.get(tid, 1)
        if rng.random() < 0.3:
            ver += 1
            version_clock[tid] = ver
        final_md5[tid] = _gen_md5(tid, ver)

    # ---- pass 2 — emit manifest entries with shared md5 ----
    direct = [
        _make_table_entry(tid, name=tid.replace("tbl_", ""), md5=final_md5[tid])
        for tid in direct_tids
    ]

    packages: List[Dict[str, Any]] = []
    for slug, ptables in pkg_specs:
        tables_payload = [
            _make_table_entry(tid, name=tid.replace("tbl_", ""), md5=final_md5[tid])
            for tid in ptables
        ]
        packages.append({
            "id": f"id_{slug}",
            "slug": slug,
            "name": slug,
            "tables": tables_payload,
        })

    domains: List[Dict[str, Any]] = []
    for slug in ("ops", "finance", "ml"):
        if rng.random() < 0.5:
            ver = version_clock.get(f"dom_{slug}", 1)
            if rng.random() < 0.3:
                ver += 1
                version_clock[f"dom_{slug}"] = ver
            domains.append({
                "slug": slug,
                "md5": _gen_md5(f"dom_{slug}", ver),
            })

    return _make_manifest(direct=direct, packages=packages, domains=domains)


# ---------------------------------------------------------------------------
# Property test
# ---------------------------------------------------------------------------


@pytest.mark.parametrize("seed", [42, 1337, 2024, 7, 99])
def test_pull_invariants_property(tmp_path, seed):
    """Random-sequence-of-operations property test.

    For each seeded RNG: build 6 random manifests in a row, pull each in
    turn, assert invariants after every pull. Reproducible — same seed →
    same manifest sequence.
    """
    rng = random.Random(seed)
    table_pool = [f"tbl_{c}" for c in "abcdef"]
    version_clock: Dict[str, int] = {}

    # Each domain slug has a per-version body, keyed by md5.
    def _bundle_for(slug: str) -> Dict[bytes, bool]:
        return {}

    # The bundle fetcher returns the body matching the manifest's current md5.
    # For simplicity we generate fresh bytes per call — the sync layer hashes
    # by file content via the manifest md5, not by remote-hash comparison.
    def _bundle_fetcher(slug: str) -> bytes:
        ver = version_clock.get(f"dom_{slug}", 1)
        return f"BUNDLE|{slug}|v{ver}".encode()

    fetcher_md5_index: Dict[str, str] = {}

    def _fetcher(url: str, dest: Path) -> None:
        # Manifest carries explicit md5 per (id, version). Stash both so
        # the next pull's md5_of() returns the expected hash and the engine
        # doesn't loop on re-fetch.
        tid = url.rstrip("/").rsplit("/", 2)[-2]
        md5 = fetcher_md5_index.get(tid, "")
        dest.parent.mkdir(parents=True, exist_ok=True)
        body = f"FAKE|{tid}|md5={md5}".encode()
        dest.write_bytes(body)

    local_dir = tmp_path / "local"
    local_dir.mkdir()

    for step in range(6):
        manifest = _build_random_manifest(
            rng, table_pool=table_pool, version_clock=version_clock,
        )
        # Rebuild md5 index from the manifest so the fetcher knows what
        # bytes to emit for each (id, version).
        fetcher_md5_index = {}
        for t in manifest["direct_tables"]:
            fetcher_md5_index[t["id"]] = t["md5"]
        for pkg in manifest["data_packages"]:
            for t in pkg["tables"]:
                fetcher_md5_index[t["id"]] = t["md5"]

        # Use a custom md5 fn that returns the manifest-declared md5 if the
        # file's actual bytes encode it. This decouples fixture-byte md5
        # from the manifest md5 so we don't have to hash actual parquet.
        def _md5(path: Path, _index=fetcher_md5_index) -> str:
            # Re-derive: file body is f"FAKE|{tid}|md5={md5}".
            try:
                text = path.read_bytes().decode("utf-8", errors="ignore")
            except OSError:
                return ""
            for tid, expected in _index.items():
                if f"|{tid}|md5={expected}" in text:
                    return expected
            return hashlib.md5(path.read_bytes()).hexdigest()

        opts = PullStackOptions(
            manifest=manifest,
            local_dir=local_dir,
            fetcher=_fetcher,
            md5_of=_md5,
            bundle_fetcher=_bundle_fetcher,
        )
        run_stack_sync(opts)

        expected_domains = {d["slug"] for d in manifest["memory_domains"]}
        _assert_invariants(local_dir, expected_domains)


def test_pull_invariants_idempotent_repull(tmp_path):
    """Re-pulling the same manifest is a no-op for state + disk."""
    local_dir = tmp_path / "local"
    local_dir.mkdir()

    md5_index = {"tbl_a": _gen_md5("tbl_a", 1)}

    def _fetcher(url: str, dest: Path) -> None:
        tid = url.rstrip("/").rsplit("/", 2)[-2]
        dest.parent.mkdir(parents=True, exist_ok=True)
        dest.write_bytes(f"FAKE|{tid}|md5={md5_index[tid]}".encode())

    def _md5(path: Path) -> str:
        text = path.read_bytes().decode("utf-8", errors="ignore")
        for tid, expected in md5_index.items():
            if f"|{tid}|md5={expected}" in text:
                return expected
        return hashlib.md5(path.read_bytes()).hexdigest()

    def _bundle(slug: str) -> bytes:
        return f"BUNDLE|{slug}".encode()

    manifest = _make_manifest(
        direct=[_make_table_entry("tbl_a", "a", md5_index["tbl_a"])],
    )
    opts = PullStackOptions(
        manifest=manifest, local_dir=local_dir,
        fetcher=_fetcher, md5_of=_md5, bundle_fetcher=_bundle,
    )
    rep1 = run_stack_sync(opts)
    assert rep1.direct_tables.added == 1
    _assert_invariants(local_dir, expected_domains_in_stack=set())

    rep2 = run_stack_sync(opts)
    # Idempotent: no adds, no updates, no removes.
    assert rep2.direct_tables.added == 0
    assert rep2.direct_tables.updated == 0
    assert rep2.direct_tables.removed == 0
    _assert_invariants(local_dir, expected_domains_in_stack=set())


def test_pull_invariants_remove_with_overlap_preserves_shared(tmp_path):
    """Two packages share table A; dropping one package keeps A's _shared
    parquet alive (the other package still references it)."""
    local_dir = tmp_path / "local"
    local_dir.mkdir()

    md5_a = _gen_md5("tbl_a", 1)
    md5_index = {"tbl_a": md5_a}

    def _fetcher(url: str, dest: Path) -> None:
        tid = url.rstrip("/").rsplit("/", 2)[-2]
        dest.parent.mkdir(parents=True, exist_ok=True)
        dest.write_bytes(f"FAKE|{tid}|md5={md5_index[tid]}".encode())

    def _md5(path: Path) -> str:
        text = path.read_bytes().decode("utf-8", errors="ignore")
        for tid, expected in md5_index.items():
            if f"|{tid}|md5={expected}" in text:
                return expected
        return hashlib.md5(path.read_bytes()).hexdigest()

    def _bundle(slug: str) -> bytes:
        return f"BUNDLE|{slug}".encode()

    # Step 1: both packages reference tbl_a.
    m1 = _make_manifest(
        packages=[
            {"id": "id_p1", "slug": "p1", "name": "P1",
             "tables": [_make_table_entry("tbl_a", "a", md5_a)]},
            {"id": "id_p2", "slug": "p2", "name": "P2",
             "tables": [_make_table_entry("tbl_a", "a", md5_a)]},
        ],
    )
    opts = PullStackOptions(
        manifest=m1, local_dir=local_dir,
        fetcher=_fetcher, md5_of=_md5, bundle_fetcher=_bundle,
    )
    run_stack_sync(opts)
    _assert_invariants(local_dir, expected_domains_in_stack=set())
    # Verify the canonical _shared file exists and exactly two refs point at it.
    shared = local_dir / "data" / "_shared" / "tbl_a.parquet"
    assert shared.exists()

    # Step 2: drop p2 → tbl_a stays (still referenced by p1).
    m2 = _make_manifest(
        packages=[
            {"id": "id_p1", "slug": "p1", "name": "P1",
             "tables": [_make_table_entry("tbl_a", "a", md5_a)]},
        ],
    )
    opts.manifest = m2
    run_stack_sync(opts)
    assert shared.exists(), "shared parquet must survive while p1 still references it"
    _assert_invariants(local_dir, expected_domains_in_stack=set())

    # Step 3: drop p1 → tbl_a now has zero refs → _shared file removed.
    m3 = _make_manifest(packages=[])
    opts.manifest = m3
    run_stack_sync(opts)
    assert not shared.exists(), "shared parquet must be deleted when ref count hits 0"
    _assert_invariants(local_dir, expected_domains_in_stack=set())


def test_pull_invariants_memory_domain_lifecycle(tmp_path):
    """Memory domain appears → bundle.md exists; domain leaves → bundle.md gone."""
    local_dir = tmp_path / "local"
    local_dir.mkdir()

    def _fetcher(url: str, dest: Path) -> None:
        dest.parent.mkdir(parents=True, exist_ok=True)
        dest.write_bytes(b"")

    def _md5(path: Path) -> str:
        return ""

    def _bundle(slug: str) -> bytes:
        return f"BUNDLE|{slug}".encode()

    md5_finance = _gen_md5("dom_finance", 1)
    m1 = _make_manifest(
        domains=[{"slug": "finance", "md5": md5_finance}],
    )
    opts = PullStackOptions(
        manifest=m1, local_dir=local_dir,
        fetcher=_fetcher, md5_of=_md5, bundle_fetcher=_bundle,
    )
    run_stack_sync(opts)
    bundle = local_dir / "memory" / "finance" / "bundle.md"
    assert bundle.exists()
    _assert_invariants(local_dir, expected_domains_in_stack={"finance"})

    # Domain leaves stack → bundle.md removed.
    m2 = _make_manifest(domains=[])
    opts.manifest = m2
    run_stack_sync(opts)
    assert not bundle.exists()
    _assert_invariants(local_dir, expected_domains_in_stack=set())