arborist/tests/test_snapshot.py

"""Tests for corpus-level snapshots.

Snapshots pin a forest of document_roots into a single Merkle root
that's bit-identical across machines that ingested the same dump.
"""

from __future__ import annotations

import time

import pytest

from aborist.document import Document
from aborist.ingest import ingest_source
from aborist.snapshot import (
    EMPTY_SNAPSHOT_ROOT,
    compute_snapshot_root,
    create_snapshot,
    diff_against_current,
    list_snapshots,
    verify_snapshot,
)
from aborist.source import Source
from aborist.store import connect


class _FakeSource(Source):
    source_type = "test"

    def __init__(self, docs):
        self._docs = docs

    def iter_documents(self):
        for d in self._docs:
            yield d


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


# ---------------------------------------------------------------------------
# Pure compute
# ---------------------------------------------------------------------------


def test_compute_snapshot_root_empty_corpus(tmp_path):
    conn = connect(tmp_path / "a.db")
    try:
        root, count = compute_snapshot_root(conn)
        assert root == EMPTY_SNAPSHOT_ROOT
        assert count == 0
    finally:
        conn.close()


def test_compute_snapshot_root_single_doc_uses_doc_root(tmp_path):
    conn = connect(tmp_path / "a.db")
    try:
        ingest_source(conn, _FakeSource([_doc("test://a", "hello world a")]))
        root, count = compute_snapshot_root(conn)
        # Degenerate single-leaf: snapshot_root == document_root.
        doc_row = conn.execute(
            "SELECT document_root FROM documents"
        ).fetchone()
        assert root == doc_row["document_root"]
        assert count == 1
    finally:
        conn.close()


def test_compute_snapshot_root_is_order_independent(tmp_path):
    """Two corpora with the same documents in different ingest order
    must produce the same snapshot_root. That's the cross-machine
    convergence property."""
    a = tmp_path / "a.db"
    b = tmp_path / "b.db"

    docs1 = [_doc("test://x", "x" * 100), _doc("test://y", "y" * 100)]
    docs2 = list(reversed(docs1))

    conn_a = connect(a)
    conn_b = connect(b)
    try:
        ingest_source(conn_a, _FakeSource(docs1))
        ingest_source(conn_b, _FakeSource(docs2))
        root_a, _ = compute_snapshot_root(conn_a)
        root_b, _ = compute_snapshot_root(conn_b)
        assert root_a == root_b
    finally:
        conn_a.close()
        conn_b.close()


def test_compute_snapshot_root_changes_when_doc_added(tmp_path):
    conn = connect(tmp_path / "a.db")
    try:
        ingest_source(conn, _FakeSource([_doc("test://a", "alpha alpha alpha alpha")]))
        root1, _ = compute_snapshot_root(conn)
        ingest_source(conn, _FakeSource([_doc("test://b", "beta beta beta beta")]))
        root2, _ = compute_snapshot_root(conn)
        assert root1 != root2
    finally:
        conn.close()


# ---------------------------------------------------------------------------
# Persistence + audit chain
# ---------------------------------------------------------------------------


def test_create_snapshot_persists_and_audits(tmp_path):
    conn = connect(tmp_path / "a.db")
    try:
        ingest_source(conn, _FakeSource([_doc("test://a", "hello world a")]))
        result = create_snapshot(conn, reason="post-ingest")
        assert result["doc_count"] == 1
        assert len(result["snapshot_root"]) == 64
        assert len(result["audit_event_hash"]) == 64

        rows = list_snapshots(conn)
        assert len(rows) == 1
        assert rows[0]["snapshot_root"] == result["snapshot_root"]
        assert rows[0]["reason"] == "post-ingest"

        # Audit chain has a snapshot_create event linked to the snapshot_root.
        audit = conn.execute(
            "SELECT event_type, subject_root FROM audit_events "
            "WHERE event_type = 'snapshot_create' ORDER BY seq DESC LIMIT 1"
        ).fetchone()
        assert audit["event_type"] == "snapshot_create"
        assert audit["subject_root"] == result["snapshot_root"]
    finally:
        conn.close()


def test_create_snapshot_auto_links_to_prior(tmp_path):
    conn = connect(tmp_path / "a.db")
    try:
        ingest_source(conn, _FakeSource([_doc("test://a", "hello world a")]))
        first = create_snapshot(conn, reason="first")
        # Add a doc; second snapshot's parent should auto-link to first.
        ingest_source(conn, _FakeSource([_doc("test://b", "beta beta beta beta")]))
        # Sleep briefly so taken_at differs even on fast machines.
        time.sleep(1.1)
        second = create_snapshot(conn, reason="second")
        assert second["snapshot_root"] != first["snapshot_root"]
        assert second["parent_snapshot"] == first["snapshot_root"]
    finally:
        conn.close()


def test_create_snapshot_idempotent_on_unchanged_corpus(tmp_path):
    conn = connect(tmp_path / "a.db")
    try:
        ingest_source(conn, _FakeSource([_doc("test://a", "hello world a")]))
        a = create_snapshot(conn, reason="first")
        b = create_snapshot(conn, reason="second-no-change")
        # Same corpus -> same snapshot_root. Second insert is a no-op
        # against the PK; we don't get a duplicate row.
        assert a["snapshot_root"] == b["snapshot_root"]
        assert len(list_snapshots(conn)) == 1
    finally:
        conn.close()


# ---------------------------------------------------------------------------
# Verify + diff
# ---------------------------------------------------------------------------


def test_verify_snapshot_matches_when_corpus_unchanged(tmp_path):
    conn = connect(tmp_path / "a.db")
    try:
        ingest_source(conn, _FakeSource([_doc("test://a", "hello world a")]))
        s = create_snapshot(conn)
        result = verify_snapshot(conn, s["snapshot_root"])
        assert result["matches"]
    finally:
        conn.close()


def test_verify_snapshot_rejects_after_drift(tmp_path):
    conn = connect(tmp_path / "a.db")
    try:
        ingest_source(conn, _FakeSource([_doc("test://a", "hello world a")]))
        s = create_snapshot(conn)
        ingest_source(conn, _FakeSource([_doc("test://b", "beta beta beta beta")]))
        result = verify_snapshot(conn, s["snapshot_root"])
        assert not result["matches"]
        assert result["current_doc_count"] == 2
    finally:
        conn.close()


def test_diff_returns_identical_for_matching_root(tmp_path):
    conn = connect(tmp_path / "a.db")
    try:
        ingest_source(conn, _FakeSource([_doc("test://a", "hello world a")]))
        s = create_snapshot(conn)
        result = diff_against_current(conn, s["snapshot_root"])
        assert result["diff"] == "identical"
    finally:
        conn.close()


def test_diff_reports_drift_after_adds(tmp_path):
    conn = connect(tmp_path / "a.db")
    try:
        ingest_source(conn, _FakeSource([_doc("test://a", "hello world a")]))
        s = create_snapshot(conn)
        ingest_source(conn, _FakeSource([_doc("test://b", "beta beta beta beta")]))
        result = diff_against_current(conn, s["snapshot_root"])
        assert result["diff"] == "drifted"
        assert result["snapshot_root"] != result["current_root"]
    finally:
        conn.close()