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_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