arborist/tests/test_cold_object.py

"""Cold-pack distribution tier (#000061) — pack build, push, pull.

Uses the in-process `MemoryBackend` so the default test suite runs without
boto3, moto, or a live S3. The on-the-wire boto3+moto test lives in
`tests/test_cold_object_boto3.py` (gated on those extras).
"""

from __future__ import annotations

import json
from typing import Iterator

import pytest

from arborist.cold_object import (
    PACK_PREFIX,
    MemoryBackend,
    build_pack,
    open_pack,
    pack_key,
    parse_manifest,
)
from arborist.compress import unpack_chunk
from arborist.document import Document
from arborist.evict import (
    pull_pack,
    push_pack,
)
from arborist.ingest import ingest_source
from arborist.merkle import hash_leaf
from arborist.source import Source
from arborist.store import connect


class FakeSource(Source):
    source_type = "html"

    def __init__(self, docs: list[Document]):
        self.docs = docs

    def iter_documents(self) -> Iterator[Document]:
        yield from self.docs


def _doc(uri: str, content: str) -> Document:
    return Document(uri=uri, content=content, source_type="html", title=uri)


LONG = (
    "The eight forms of capital include living, social, and intellectual. " * 30
    + "\n\n"
    + "Merkle providence proves answer derives from a specific source. " * 30
)


# ---------------------------------------------------------------------------
# Pure-helper tests (no DB, no backend) — the content-addressing invariants
# ---------------------------------------------------------------------------


def test_pack_key_layout():
    h = "b" * 64
    assert pack_key(h).endswith(".tar.zst")
    assert pack_key(h, manifest=True).endswith(".manifest.ndjson")
    # Different scheme so a `list_keys("packs/")` can distinguish.
    assert pack_key(h) != pack_key(h, manifest=True)


def test_build_pack_is_content_addressed():
    # Same input → same pack_hash, regardless of ordering.
    body_a = b"hello world " * 100
    body_b = b"goodbye world " * 100
    leaf_a = hash_leaf(body_a).hex()
    leaf_b = hash_leaf(body_b).hex()

    pack1 = build_pack([(leaf_a, body_a), (leaf_b, body_b)])
    pack2 = build_pack([(leaf_b, body_b), (leaf_a, body_a)])  # reversed
    assert pack1.pack_hash == pack2.pack_hash
    assert {e.leaf_hash for e in pack1.entries} == {leaf_a, leaf_b}


def test_build_pack_rejects_empty():
    with pytest.raises(ValueError):
        build_pack([])


def test_open_pack_verifies_each_chunk_hash():
    body_a = b"chunk one body padded out " * 20
    body_b = b"chunk two body padded out " * 20
    leaf_a = hash_leaf(body_a).hex()
    leaf_b = hash_leaf(body_b).hex()
    pack = build_pack([(leaf_a, body_a), (leaf_b, body_b)])

    recovered = dict(open_pack(pack.body_bytes))
    assert recovered[leaf_a] == body_a
    assert recovered[leaf_b] == body_b


def test_parse_manifest_round_trip():
    body = b"abc" * 100
    leaf = hash_leaf(body).hex()
    pack = build_pack([(leaf, body)])
    entries = parse_manifest(pack.manifest_bytes)
    assert len(entries) == 1
    assert entries[0].leaf_hash == leaf
    assert entries[0].size == len(body)


# ---------------------------------------------------------------------------
# Backend round-trip — MemoryBackend, no network
# ---------------------------------------------------------------------------


def test_memory_backend_put_get_head_list():
    """The backend speaks raw bytes; pack-level semantics live above it."""
    b = MemoryBackend()
    key = "packs/abc123.tar.zst"
    payload = b"fake pack body " * 10
    assert not b.head(key)
    b.put(key, payload, content_type="application/zstd")
    assert b.head(key)
    assert b.get(key) == payload
    assert list(b.list_keys("packs/")) == [key]


def test_memory_backend_object_size_returns_content_length_or_none():
    b = MemoryBackend()
    key = "packs/abc.tar.zst"
    assert b.object_size(key) is None  # missing
    b.put(key, b"x" * 1234)
    assert b.object_size(key) == 1234
    assert b.object_size("packs/nonexistent") is None


def test_memory_backend_identity_carries_no_credentials():
    b = MemoryBackend(endpoint_url="memory://nyc3", bucket="arborist-test")
    body = b.identity.to_audit_body()
    # Operation Voyeur: no credential keys should ever appear.
    text = json.dumps(body)
    for k in ("access", "secret", "key", "token", "password"):
        assert k not in text.lower() or k == "key" and "access_key" not in text.lower()


# ---------------------------------------------------------------------------
# Pack push/pull — the only path now (no individual blobs)
# ---------------------------------------------------------------------------


def test_push_pack_then_pull_pack_restores_chunks(tmp_path):
    db = tmp_path / "pack.db"
    conn = connect(db)
    backend = MemoryBackend()
    try:
        ingest_source(conn, FakeSource([_doc("html://p", LONG)]))
        root = conn.execute(
            "SELECT document_root FROM documents WHERE document_uri='html://p'"
        ).fetchone()["document_root"]

        # Pack while content is still local. Tiny corpus → one pack.
        push_result = push_pack(conn, backend, document_root=root)
        assert push_result["status"] == "pushed"
        assert push_result["pack_count"] == 1
        # Pack is bound to a snapshot_root — packs are delayed point-in-time
        # snapshots; the audit row pins which corpus state this pack covers
        # so falsifications between repacks produce a different pack_hash.
        assert len(push_result["snapshot_root"]) == 64
        first = push_result["packs"][0]
        assert first["chunk_count"] > 0
        assert first["compressed_bytes"] <= first["uncompressed_bytes"] + 200
        assert first["snapshot_root"] == push_result["snapshot_root"]
        pack_hash = first["pack_hash"]
        assert backend.head(pack_key(pack_hash))
        assert backend.head(pack_key(pack_hash, manifest=True))

        # Now NULL local content so pull_pack has something to do.
        conn.execute("UPDATE chunks SET content = NULL, tier='cold'")
        conn.execute("DELETE FROM chunks_fts")
        conn.commit()

        pull_result = pull_pack(conn, backend, pack_hash)
        assert pull_result["status"] == "pulled"
        assert pull_result["chunks_restored"] == first["chunk_count"]

        cold = conn.execute(
            "SELECT COUNT(*) FROM chunks WHERE tier='cold'"
        ).fetchone()[0]
        assert cold == 0
    finally:
        conn.close()


def test_push_pack_idempotent(tmp_path):
    """Same chunk set → same pack_hash(es) → bucket overwrite is a no-op."""
    db = tmp_path / "pack-id.db"
    conn = connect(db)
    backend = MemoryBackend()
    try:
        ingest_source(conn, FakeSource([_doc("html://p", LONG)]))
        a = push_pack(conn, backend)
        b = push_pack(conn, backend)
        assert [p["pack_hash"] for p in a["packs"]] == [
            p["pack_hash"] for p in b["packs"]
        ]
    finally:
        conn.close()


def test_push_pack_nothing_to_pack(tmp_path):
    db = tmp_path / "empty.db"
    conn = connect(db)
    backend = MemoryBackend()
    try:
        result = push_pack(conn, backend)
        assert result["status"] == "nothing_to_pack"
        assert result["packs"] == []
    finally:
        conn.close()


def test_push_pack_splits_to_fit_dvdr(tmp_path):
    """A small per-pack COMPRESSED-bytes cap forces multi-pack splitting;
    every pack's compressed body stays ≤ cap + tar-trailer slack.

    The cap is on compressed bytes (`stream_packs` peeks the compressed
    buffer after each chunk via FLUSH_BLOCK). Each pack's compressed_bytes
    should be close to the cap — fill the disc, don't leave 50% empty.
    """
    db = tmp_path / "split.db"
    conn = connect(db)
    backend = MemoryBackend()
    try:
        # Distinct content per doc so zstd can't dedupe across them — the
        # whole point of testing the compressed cap is that compressed
        # output scales with corpus, not with one-pattern repetition.
        import random
        rng = random.Random(42)
        words = ["alpha", "bravo", "charlie", "delta", "echo", "foxtrot",
                 "golf", "hotel", "india", "juliet", "kilo", "lima"]
        docs = []
        for i in range(8):
            phrases = [" ".join(rng.choices(words, k=10)) for _ in range(80)]
            docs.append(_doc(f"html://d{i}", " ".join(phrases) + ". "))
        ingest_source(conn, FakeSource(docs))

        # Small compressed cap (4 KB) — should produce multiple packs.
        CAP = 4_096
        result = push_pack(conn, backend, max_pack_bytes=CAP)
        assert result["status"] == "pushed"
        assert result["pack_count"] > 1, (
            f"expected multiple packs at {CAP}-byte compressed cap, "
            f"got {result['pack_count']}"
        )

        # Tar trailer (~1 KB padding) + zstd frame footer get emitted after
        # the last in-loop size check, so overshoot is bounded by trailer
        # size + one chunk's worth of compressed bytes.
        OVERSHOOT_SLACK = 4_096
        for p in result["packs"]:
            assert p["compressed_bytes"] <= CAP + OVERSHOOT_SLACK, (
                f"pack {p['pack_hash'][:8]} compressed={p['compressed_bytes']} "
                f"exceeds cap {CAP}+slack {OVERSHOOT_SLACK}"
            )
            assert p["chunk_count"] >= 1

        # Round-trip: NULL local content, pull every pack back.
        conn.execute("UPDATE chunks SET content = NULL, tier='cold'")
        conn.execute("DELETE FROM chunks_fts")
        conn.commit()

        total_restored = 0
        for p in result["packs"]:
            r = pull_pack(conn, backend, p["pack_hash"])
            total_restored += r["chunks_restored"]
        assert total_restored == result["total_chunks"]
    finally:
        conn.close()


def test_push_pack_local_dir_writes_files(tmp_path):
    """--local-dir writes pack + manifest files for burning to physical media."""
    db = tmp_path / "burn.db"
    conn = connect(db)
    backend = MemoryBackend()
    burn_dir = tmp_path / "discs"
    try:
        ingest_source(conn, FakeSource([_doc("html://burn", LONG)]))
        result = push_pack(conn, backend, local_dir=str(burn_dir))
        assert result["status"] == "pushed"
        assert result["local_dir"] == str(burn_dir)
        assert burn_dir.exists()
        pack_files = sorted(burn_dir.glob("arborist-pack-*.tar.zst"))
        manifest_files = sorted(burn_dir.glob("arborist-pack-*.manifest.ndjson"))
        assert len(pack_files) == result["pack_count"]
        assert len(manifest_files) == result["pack_count"]
        # Local files are byte-identical to what we'd burn.
        for pack_file, summary in zip(pack_files, result["packs"]):
            assert pack_file.stat().st_size == summary["compressed_bytes"]
    finally:
        conn.close()


def test_edges_fan_in_batches_huge_destinations(tmp_path):
    """Regression test for the memory-bounded edges dump path.

    A destination with > _FAN_IN_BATCH inbound links should be split
    into multiple JSON rows on dump, and restored to the live schema
    with no missing or duplicated edges. Without this batching the
    groupby accumulator held the whole src_roots list in Python heap
    (~5 GB per worker observed on the v2 corpus run).
    """
    import sqlite3
    from arborist.cold_pack_metadata import (
        _dump_edges_fan_in,
        _restore_edges_fan_out,
        _FAN_IN_BATCH,
    )
    from arborist.store import connect as store_connect

    # Build a minimal schema-bearing DB and stuff edges directly.
    src_db = tmp_path / "edges.db"
    conn = store_connect(src_db)
    try:
        # Insert N_LINKS edges all pointing at the same dst_uri so the
        # groupby for this destination spans multiple batches.
        N_LINKS = _FAN_IN_BATCH + 137  # one full batch + a partial
        dst_uri = "https://example.com/popular"
        edge_type = "wikilink"
        anchor = ""
        dst_root = ""
        with conn:
            conn.executemany(
                "INSERT OR IGNORE INTO edges "
                "(src_root, dst_root, dst_uri, edge_type, anchor) "
                "VALUES (?, ?, ?, ?, ?)",
                [
                    (f"src{i:08d}" + "0" * 24, dst_root, dst_uri, edge_type, anchor)
                    for i in range(N_LINKS)
                ],
            )

        # Dump
        dump_path = tmp_path / "edges.jsonl"
        n_rows_written = _dump_edges_fan_in(conn, dump_path)
        # One destination split into ceil(N_LINKS / FAN_IN_BATCH) rows.
        expected_rows = (N_LINKS + _FAN_IN_BATCH - 1) // _FAN_IN_BATCH
        assert n_rows_written == expected_rows, (
            f"expected {expected_rows} fan-in rows for {N_LINKS} edges, "
            f"got {n_rows_written}"
        )

        # Wipe edges and restore. INSERT OR IGNORE handles any duplicate
        # PK rows from chunk boundaries cleanly.
        with conn:
            conn.execute("DELETE FROM edges")
        n_restored = _restore_edges_fan_out(conn, dump_path)
        conn.commit()
        assert n_restored == N_LINKS, (
            f"expected {N_LINKS} edges restored, got {n_restored}"
        )

        # Verify the live table has exactly N_LINKS rows for that destination.
        final_count = conn.execute(
            "SELECT COUNT(*) FROM edges WHERE dst_uri = ?",
            (dst_uri,),
        ).fetchone()[0]
        assert final_count == N_LINKS
    finally:
        conn.close()


def test_push_pack_v2_hydrates_fresh_empty_db(tmp_path):
    """Pack format v2 self-sufficiency test: build a pack from a populated
    DB, then unpack it into a FRESH empty DB and verify every table is
    restored (not just chunks content). This is the real new-peer
    hydration path."""
    src_db = tmp_path / "src.db"
    rx_db = tmp_path / "rx.db"
    conn_src = connect(src_db)
    backend = MemoryBackend()
    try:
        # Populate the source DB with a few docs so multiple tables get
        # rows: documents, chunks, merkle_nodes, audit_events at minimum.
        # Distinct content per doc — same content dedupes by document_root.
        docs = [
            _doc(f"html://hydrate{i}", LONG + f"\n\nDoc-specific marker {i}." * 5)
            for i in range(3)
        ]
        ingest_source(conn_src, FakeSource(docs))

        src_docs = conn_src.execute(
            "SELECT COUNT(*) FROM documents"
        ).fetchone()[0]
        src_chunks = conn_src.execute(
            "SELECT COUNT(*) FROM chunks"
        ).fetchone()[0]
        src_merkle = conn_src.execute(
            "SELECT COUNT(*) FROM merkle_nodes"
        ).fetchone()[0]
        src_audit = conn_src.execute(
            "SELECT COUNT(*) FROM audit_events"
        ).fetchone()[0]
        assert src_docs == 3 and src_chunks > 0 and src_merkle > 0 and src_audit > 0

        result = push_pack(conn_src, backend)
        assert result["status"] == "pushed"
        pack_hash = result["packs"][0]["pack_hash"]
    finally:
        conn_src.close()

    # Brand-new empty receiving DB. No ingest, nothing — just `connect()`
    # which creates the schema.
    conn_rx = connect(rx_db)
    try:
        # Confirm rx is empty before unpack.
        assert conn_rx.execute(
            "SELECT COUNT(*) FROM documents"
        ).fetchone()[0] == 0

        pull_result = pull_pack(conn_rx, backend, pack_hash)
        assert pull_result["status"] == "pulled"
        assert pull_result["pack_format"] == "v2"
        # Every table that was non-empty in src should have rows now in rx.
        assert pull_result["tables_restored"]["documents"] == src_docs
        assert pull_result["tables_restored"]["chunks"] == src_chunks
        assert pull_result["tables_restored"]["merkle_nodes"] == src_merkle
        assert pull_result["tables_restored"]["audit_events"] == src_audit
        # Chunk bodies restored.
        rx_hot = conn_rx.execute(
            "SELECT COUNT(*) FROM chunks WHERE tier='hot' AND content IS NOT NULL"
        ).fetchone()[0]
        assert rx_hot == src_chunks
        # FTS5 rows rebuilt.
        rx_fts = conn_rx.execute("SELECT COUNT(*) FROM chunks_fts").fetchone()[0]
        assert rx_fts == src_chunks
        # All three documents reachable by URI.
        for i in range(3):
            row = conn_rx.execute(
                "SELECT title FROM documents WHERE document_uri=?",
                (f"html://hydrate{i}",),
            ).fetchone()
            assert row is not None, f"document html://hydrate{i} not restored"
    finally:
        conn_rx.close()


def test_push_pack_no_push_writes_locally_only(tmp_path):
    """push_to_bucket=False skips bucket writes, requires local_dir."""
    db = tmp_path / "local-only.db"
    conn = connect(db)
    backend = MemoryBackend()
    burn_dir = tmp_path / "iso"
    try:
        ingest_source(conn, FakeSource([_doc("html://x", LONG)]))
        result = push_pack(
            conn, backend,
            local_dir=str(burn_dir),
            push_to_bucket=False,
        )
        assert result["status"] == "written"
        # No keys in bucket.
        assert list(backend.list_keys()) == []
        # Files on disk.
        assert any(burn_dir.glob("arborist-pack-*.tar.zst"))
    finally:
        conn.close()


def test_push_pack_is_deterministic_across_runs(tmp_path):
    """Same chunk set on two runs against the same DB produces the same
    pack_hashes — `ORDER BY leaf_hash` makes the chunk-to-pack assignment
    a function of (chunk set, cap) and nothing else. This is the
    prerequisite for parallel pack workers and for two replicas at the
    same snapshot to produce byte-identical bucket state."""
    db = tmp_path / "det.db"
    conn = connect(db)
    try:
        # Use distinct content so multiple chunks survive dedupe and
        # we have enough material to potentially split.
        import random
        rng = random.Random(7)
        words = ["alpha", "bravo", "charlie", "delta", "echo", "foxtrot"]
        docs = []
        for i in range(6):
            phrases = [" ".join(rng.choices(words, k=10)) for _ in range(80)]
            docs.append(_doc(f"html://d{i}", " ".join(phrases) + ". "))
        ingest_source(conn, FakeSource(docs))

        backend_a = MemoryBackend()
        backend_b = MemoryBackend()
        result_a = push_pack(conn, backend_a, max_pack_bytes=4_000)
        result_b = push_pack(conn, backend_b, max_pack_bytes=4_000)

        hashes_a = [p["pack_hash"] for p in result_a["packs"]]
        hashes_b = [p["pack_hash"] for p in result_b["packs"]]
        assert hashes_a == hashes_b, (
            "two runs at the same snapshot produced different pack_hashes "
            f"(a={hashes_a} vs b={hashes_b}) — ORDER BY drift?"
        )
    finally:
        conn.close()


def test_default_max_pack_bytes_is_dvdr_safe_fit():
    """Default cap is 4.4 GB — DVD-R safe-fit, with ~6.5% buffer below the
    4.7 GB media spec to absorb ISO9660 overhead, media manufacturing
    variance, and drive-edge refusal. Don't drift this without a deliberate
    decision; downstream burning workflows assume it."""
    from arborist.evict import DEFAULT_MAX_PACK_BYTES
    assert DEFAULT_MAX_PACK_BYTES == 4_400_000_000
    # Must sit well below the 4.7 GB marketing capacity AND below the
    # 4,706,074,624 ECMA-267 physical capacity, with room for overhead.
    assert DEFAULT_MAX_PACK_BYTES < 4_700_000_000
    assert (4_700_000_000 - DEFAULT_MAX_PACK_BYTES) >= 250_000_000  # >=250 MB headroom