arborist/tests/test_ingest.py

"""End-to-end ingest test using a hand-rolled in-memory Source."""

from __future__ import annotations

from typing import Iterator

from arborist.document import Document, Edge
from arborist.ingest import ingest_source, verify_random_sample
from arborist.search import FTS5Backend
from arborist.search.base import AuditMode
from arborist.source import Source
from arborist.store import connect, stats


class FakeSource(Source):
    source_type = "fake"

    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, *, edges: list[Edge] | None = None) -> Document:
    return Document(
        uri=uri,
        content=content,
        source_type="fake",
        title=uri.rsplit("/", 1)[-1],
        edges=edges or [],
    )


def test_ingest_basic_round_trip(tmp_path):
    db_path = tmp_path / "test.db"
    src = FakeSource([
        _doc("test://a", "alpha bravo charlie delta echo foxtrot golf hotel"),
        _doc("test://b", "the quick brown fox jumps over the lazy dog"),
        _doc(
            "test://c",
            "merkle providence reverse rag verifies provenance",
            edges=[Edge(edge_type="wikilink", dst_uri="test://a")],
        ),
    ])
    conn = connect(db_path)
    try:
        result = ingest_source(conn, src)
        assert result.seen == 3
        assert result.inserted == 3
        assert result.skipped_duplicate == 0

        # Verify Merkle round-trip.
        v = verify_random_sample(conn, n=3)
        assert v["sampled"] == 3
        assert v["passed"] == 3
        assert v["failed"] == 0

        # Idempotent re-ingest.
        result2 = ingest_source(conn, src)
        assert result2.inserted == 0
        assert result2.skipped_duplicate == 3

        # FTS5 search returns UNGROUNDED hits.
        backend = FTS5Backend(conn)
        hits = backend.search("merkle")
        assert len(hits) >= 1
        assert hits[0].audit_mode == AuditMode.UNGROUNDED
        assert "merkle" in hits[0].snippet.lower()

        # Edge resolution: c -> a should be backfilled (a was ingested first).
        row = conn.execute(
            "SELECT dst_root FROM edges WHERE dst_uri = ?", ("test://a",)
        ).fetchone()
        assert row is not None
        assert row["dst_root"] != ""  # backfilled (was '' before resolution)

        # Stats reflect ingest.
        s = stats(conn)
        assert s["documents_total"] == 3
        assert s["documents_surface"] == 3
        assert s["documents_core"] == 0
        assert s["chunks_total"] >= 3
        assert s["audit_events_total"] == 3  # one ingest event per doc
    finally:
        conn.close()


def test_audit_chain_links_correctly(tmp_path):
    """Each audit event chains to the previous via prev_event_hash."""
    db_path = tmp_path / "audit.db"
    src = FakeSource([_doc(f"test://{i}", f"document number {i} content") for i in range(5)])
    conn = connect(db_path)
    try:
        ingest_source(conn, src)
        rows = conn.execute(
            "SELECT seq, event_hash, prev_event_hash FROM audit_events ORDER BY seq"
        ).fetchall()
        assert len(rows) == 5
        assert rows[0]["prev_event_hash"] is None  # genesis
        for i in range(1, len(rows)):
            assert rows[i]["prev_event_hash"] == rows[i - 1]["event_hash"]
    finally:
        conn.close()


def test_chunker_version_persisted(tmp_path):
    db_path = tmp_path / "chunker.db"
    conn = connect(db_path)
    try:
        ingest_source(conn, FakeSource([_doc("test://x", "alpha beta gamma")]))
        row = conn.execute(
            "SELECT chunking_version, canonicalization_version, schema_version FROM documents"
        ).fetchone()
        assert row["chunking_version"] == "tok-512-v1"
        assert row["canonicalization_version"] == "norm-v1"
        assert row["schema_version"] == "v9.8.0"
    finally:
        conn.close()