arborist/tests/test_qa_corpus.py

"""Unit tests for arborist.qa.corpus — adapter contract conformance.

What this asserts:
  * Both adapters implement the Corpus Protocol (runtime_checkable)
  * Hit + ChunkRow are frozen and hashable
  * NotSupportedError raised by routes adapters don't implement
  * SqliteShardCorpus.fts_body sanitizes natural-language input
    (no FTS5 syntax leaks through punctuation)
  * Empty / stopword-only queries return [] rather than raising

What this does NOT assert (covered in integration / functional):
  * Cross-adapter result agreement
  * Per-chunk content decompression equivalence
  * End-to-end LLM + verifier behaviour
"""
from __future__ import annotations

import sqlite3
from typing import Iterator

import pytest

from arborist.document import Document
from arborist.ingest import ingest_source
from arborist.qa.corpus import (
    ChunkRow,
    Corpus,
    Hit,
    NotSupportedError,
    SidecarBucketCorpus,
    SqliteShardCorpus,
)
from arborist.source import Source
from arborist.store import connect


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------


class _FakeSource(Source):
    source_type = "test"
    def __init__(self, ds): self.ds = ds
    def iter_documents(self) -> Iterator[Document]:
        yield from self.ds


def _doc(uri: str, content: str, title: str | None = None) -> Document:
    return Document(
        uri=uri, content=content, source_type="test",
        title=title or uri.split("/")[-1],
    )


CORPUS_DOCS = [
    _doc("test://anarchism", (
        "Anarchism is a political philosophy that promotes a stateless society. "
        * 6
        + "It seeks to abolish authority. " * 6
    ), title="Anarchism"),
    _doc("test://capital", (
        "The eight forms of capital include living, social, intellectual. " * 6
    ), title="Capital"),
]


@pytest.fixture
def shard_conn(tmp_path):
    db = tmp_path / "corpus.db"
    c = connect(db)
    ingest_source(c, _FakeSource(CORPUS_DOCS))
    yield c
    c.close()


# ---------------------------------------------------------------------------
# Data-class invariants
# ---------------------------------------------------------------------------


def test_hit_is_mutable_and_equals_by_value():
    """Hit is mutable on purpose so the rerank pipeline in
    arborist.qa.retrieval_routes can mutate .score in place (matches
    legacy _Hit convention). Equality stays value-based."""
    h1 = Hit("root1", "uri1", "title1", 1.5)
    h2 = Hit("root1", "uri1", "title1", 1.5)
    assert h1 == h2             # value equality
    # Mutation works — score is the lever rerank stages use.
    h1.score = 2.0
    assert h1.score == 2.0
    assert h1 != h2             # mutation breaks equality (as expected)
    # extras carries a dict so Hit isn't hashable by design — callers
    # dedup by (document_root, ...) tuples, not by put-in-a-set.


def test_chunk_row_is_frozen():
    c = ChunkRow("root", 0, "hex", "content")
    with pytest.raises(Exception):
        c.idx = 1                # type: ignore[misc]


def test_hit_extras_defaults_empty():
    h = Hit("a", "b", "c", 0.0)
    assert h.extras == {}


# ---------------------------------------------------------------------------
# Protocol conformance
# ---------------------------------------------------------------------------


def test_sqlite_shard_corpus_satisfies_protocol(shard_conn):
    corpus = SqliteShardCorpus(shard_conn)
    assert isinstance(corpus, Corpus)
    assert corpus.name == "sqlite-shard"
    assert corpus.higher_is_better is False


def test_sidecar_bucket_corpus_satisfies_protocol():
    # Doesn't need a live bucket — runtime_checkable Protocol only
    # checks method presence, not call signatures.
    class _StubMulti:
        manifest = type("M", (), {"shards": [], "snapshot_root": None})()
        def fts_search(self, q, *, limit=8): return []
        def conn_for_shard(self, url): raise KeyError(url)
    corpus = SidecarBucketCorpus(_StubMulti())
    assert isinstance(corpus, Corpus)
    assert corpus.name == "sidecar-bucket"
    # FTS5 bm25 is negative; lower (more negative) = better.
    assert corpus.higher_is_better is False


# ---------------------------------------------------------------------------
# fts_body sanitization (no FTS5-syntax leaks)
# ---------------------------------------------------------------------------


def test_sqlite_fts_body_handles_natural_language(shard_conn):
    """Punctuation like '?' and '-' would crash raw FTS5. The adapter
    must sanitize via _to_fts5 (mirrors cloud sanitizer)."""
    corpus = SqliteShardCorpus(shard_conn)
    # Question mark + apostrophe must not crash.
    hits = corpus.fts_body("what is anarchism?", limit=5)
    assert any("Anarchism" in h.title for h in hits)


def test_sqlite_fts_body_stopword_only_query_returns_empty(shard_conn):
    """All-stopword queries reduce to empty FTS5 expressions; the
    adapter must short-circuit to [] rather than passing '' to MATCH."""
    corpus = SqliteShardCorpus(shard_conn)
    hits = corpus.fts_body("the a an of with", limit=5)
    assert hits == []


def test_sqlite_fts_body_returns_hit_shape(shard_conn):
    corpus = SqliteShardCorpus(shard_conn)
    hits = corpus.fts_body("anarchism", limit=2)
    assert len(hits) >= 1
    for h in hits:
        assert isinstance(h, Hit)
        assert h.document_root and h.document_uri and h.title
        assert isinstance(h.score, float)


# ---------------------------------------------------------------------------
# fts_title + fts_phrase routes (documents_fts + FTS5 phrase MATCH)
# ---------------------------------------------------------------------------


def test_sqlite_fts_title_via_documents_fts(shard_conn):
    """fts_title returns title-bm25 ranked hits from documents_fts."""
    corpus = SqliteShardCorpus(shard_conn)
    hits = corpus.fts_title("anarchism", limit=5)
    # Should rank the Anarchism doc highly via title MATCH.
    assert any("Anarchism" in h.title for h in hits)


def test_sqlite_fts_title_stopword_only_returns_empty(shard_conn):
    """All-stopword queries reduce to empty FTS5 — short-circuit."""
    corpus = SqliteShardCorpus(shard_conn)
    assert corpus.fts_title("the of and") == []


def test_sqlite_fts_phrase_matches_verbatim(shard_conn):
    """FTS5 phrase MATCH '"phrase"' returns chunks containing the
    verbatim phrase. Use a phrase known to appear in the fixture doc."""
    corpus = SqliteShardCorpus(shard_conn)
    # The Anarchism fixture has "political philosophy" verbatim.
    hits = corpus.fts_phrase(["political philosophy"], limit=5)
    assert hits, "expected at least one hit for verbatim phrase"
    assert any("Anarchism" in h.title for h in hits)
    # Phrase tokens that DON'T appear together should return nothing.
    misses = corpus.fts_phrase(["stateless capital"], limit=5)
    assert misses == []


def test_sqlite_fts_phrase_empty_input_returns_empty(shard_conn):
    corpus = SqliteShardCorpus(shard_conn)
    assert corpus.fts_phrase([]) == []
    assert corpus.fts_phrase(["", "  "]) == []


def test_sidecar_fts_title_delegates_to_multishard():
    """SidecarBucketCorpus.fts_title delegates to the multi-shard
    fts_title_search method. Stub the multi-shard and confirm shape."""
    class _StubMulti:
        manifest = type("M", (), {"shards": [], "snapshot_root": None})()
        def fts_search(self, q, *, limit=8): return []
        def fts_title_search(self, q, *, limit=8): return [
            {"document_root": "abc", "document_uri": "test://x",
             "title": "X", "score": -12.0, "_shard_url": "s0"},
        ]
        def fts_phrase_search(self, ngs, *, limit=8): return []
        def conn_for_shard(self, url): raise KeyError(url)
    corpus = SidecarBucketCorpus(_StubMulti())
    hits = corpus.fts_title("anything")
    assert len(hits) == 1
    assert hits[0].title == "X"
    assert hits[0].shard_id == "s0"


def test_sidecar_fts_phrase_delegates_to_multishard():
    class _StubMulti:
        manifest = type("M", (), {"shards": [], "snapshot_root": None})()
        def fts_search(self, q, *, limit=8): return []
        def fts_title_search(self, q, *, limit=8): return []
        def fts_phrase_search(self, ngs, *, limit=8): return [
            {"document_root": "abc", "document_uri": "test://x",
             "title": "X", "score": -23.5},
        ]
        def conn_for_shard(self, url): raise KeyError(url)
    corpus = SidecarBucketCorpus(_StubMulti())
    hits = corpus.fts_phrase(["any phrase"])
    assert len(hits) == 1
    assert hits[0].score == -23.5


# ---------------------------------------------------------------------------
# chunks_for_doc shape + ordering
# ---------------------------------------------------------------------------


def test_sqlite_chunks_for_doc_returns_idx_ascending(shard_conn):
    corpus = SqliteShardCorpus(shard_conn)
    hits = corpus.fts_body("anarchism", limit=1)
    assert hits
    chunks = corpus.chunks_for_doc(hits[0].document_root)
    assert len(chunks) >= 1
    indices = [c.idx for c in chunks]
    assert indices == sorted(indices)
    for c in chunks:
        assert isinstance(c, ChunkRow)
        assert c.content       # decompressed, non-empty
        assert c.leaf_hash      # ingest sets this


def test_sqlite_snapshot_root_hex(shard_conn):
    corpus = SqliteShardCorpus(shard_conn)
    root = corpus.snapshot_root()
    assert isinstance(root, str)
    assert len(root) == 64
    int(root, 16)   # well-formed hex


# ---------------------------------------------------------------------------
# MultiShardSqliteCorpus — per-shard FTS + merge
# ---------------------------------------------------------------------------


def test_multishard_sqlite_aggregates_hits_across_shards(tmp_path):
    """Two shards each holding a distinct anchor doc. fts_body must
    surface BOTH topic-matches when their respective terms are queried,
    proving the per-shard merge works rather than the old single-conn
    rowid-namespace failure."""
    from arborist.qa.corpus import MultiShardSqliteCorpus

    # Shard A: anarchism content only.
    db_a = tmp_path / "a.db"
    c = connect(db_a)
    ingest_source(c, _FakeSource([CORPUS_DOCS[0]]))  # Anarchism
    c.close()
    # Shard B: capital content only.
    db_b = tmp_path / "b.db"
    c = connect(db_b)
    ingest_source(c, _FakeSource([CORPUS_DOCS[1]]))  # Capital
    c.close()

    corpus = MultiShardSqliteCorpus([db_a, db_b])
    try:
        # 'anarchism' hits only shard A.
        a_hits = corpus.fts_body("anarchism", limit=3)
        assert any("Anarchism" in h.title for h in a_hits)

        # 'capital' hits only shard B.
        b_hits = corpus.fts_body("capital", limit=3)
        assert any("Capital" in h.title for h in b_hits)

        # snapshot_root spans BOTH shards' document_roots.
        root = corpus.snapshot_root()
        assert len(root) == 64

        # Hit.shard_id annotated → chunks_for_doc can route.
        assert all(h.shard_id is not None for h in a_hits + b_hits)
    finally:
        corpus.close()


def test_multishard_sqlite_chunks_for_doc_finds_right_shard(tmp_path):
    """chunks_for_doc walks shards until one returns rows. Doc in
    shard B should still resolve when shard A is asked first."""
    from arborist.qa.corpus import MultiShardSqliteCorpus

    db_a = tmp_path / "a.db"
    c = connect(db_a); ingest_source(c, _FakeSource([CORPUS_DOCS[0]])); c.close()
    db_b = tmp_path / "b.db"
    c = connect(db_b); ingest_source(c, _FakeSource([CORPUS_DOCS[1]])); c.close()

    corpus = MultiShardSqliteCorpus([db_a, db_b])
    try:
        b_hits = corpus.fts_body("capital", limit=1)
        assert b_hits
        chunks = corpus.chunks_for_doc(b_hits[0].document_root)
        assert chunks, "chunks_for_doc returned nothing for shard-B-owned doc"
        assert any("capital" in c.content.lower() for c in chunks)
    finally:
        corpus.close()