arborist/tests/test_warrant_chain.py

"""Tests for Phase 3 of `#000031` — warrant-chain lookup +
verifier suppression of WARRANT_MISSING when a Merkle-bound
primary-source backing exists.
"""

from __future__ import annotations

import sqlite3
from pathlib import Path

import pytest

from arborist.qa.warrant_chain import (
    has_warrant_chain,
    warrant_chain_lookup,
)
from arborist.store import SCHEMA_SQL


@pytest.fixture
def shards_with_chain(tmp_path: Path):
    """Build a minimal two-shard cluster: a main shard with a
    `derivations` row written by warrant-resolver-v1 (mimics what
    `arborist warrant-resolve --write` produces) plus a sibling
    crawl shard that doesn't have a derivations table.

    Returns ``(shards_dir, core_root, src_root)``.
    """
    shards_dir = tmp_path / "shards"
    crawl_dir = tmp_path / "crawl"
    shards_dir.mkdir()
    crawl_dir.mkdir()

    main_db = shards_dir / "000.db"
    conn = sqlite3.connect(str(main_db))
    try:
        conn.executescript(SCHEMA_SQL)
        # Insert a fake claim-pack record + textbook surface document
        # so the FK constraints on derivations are satisfied.
        core_root = "a" * 64
        src_root = "b" * 64
        conn.execute(
            "INSERT INTO documents "
            "(document_root, document_uri, source_type, kind, "
            " compression_depth, title, chunking_version, "
            " canonicalization_version, schema_version, ingest_ts) "
            "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
            (core_root, "claim_pack://test", "claim_pack", "surface",
             0, "Test record", "tok-512-v1", "norm-v1", "v9.8.0", 0),
        )
        conn.execute(
            "INSERT INTO documents "
            "(document_root, document_uri, source_type, kind, "
            " compression_depth, title, chunking_version, "
            " canonicalization_version, schema_version, ingest_ts) "
            "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
            (src_root, "https://example/textbook", "html", "surface",
             0, "Textbook", "tok-512-v1", "norm-v1", "v9.8.0", 0),
        )
        conn.execute(
            "INSERT INTO derivations "
            "(core_root, src_root, process_id, proof_blob, distilled_at) "
            "VALUES (?, ?, ?, ?, ?)",
            (core_root, src_root, "warrant-resolver-v1", b"{}", 0),
        )
        conn.commit()
    finally:
        conn.close()

    # Sibling crawl shard with no derivations table — exercises the
    # try/except for OperationalError.
    sibling_db = crawl_dir / "textbook_test.db"
    sib = sqlite3.connect(str(sibling_db))
    try:
        sib.execute("CREATE TABLE meta (k TEXT, v TEXT)")
        sib.commit()
    finally:
        sib.close()

    return shards_dir, core_root, src_root


# --- warrant_chain_lookup -------------------------------------


def test_warrant_chain_lookup_finds_resolver_rows(shards_with_chain):
    shards_dir, core_root, _src_root = shards_with_chain
    roots = warrant_chain_lookup(shards_dir)
    assert core_root in roots
    assert isinstance(roots, frozenset)


def test_warrant_chain_lookup_skips_unrelated_process_ids(tmp_path):
    """A shard whose derivations are from non-warrant processes
    (e.g. distillation) should NOT contribute to the warrant-chain set.
    """
    shards_dir = tmp_path / "shards"
    shards_dir.mkdir()
    db = shards_dir / "000.db"
    conn = sqlite3.connect(str(db))
    try:
        conn.executescript(SCHEMA_SQL)
        conn.execute(
            "INSERT INTO documents "
            "(document_root, document_uri, source_type, kind, "
            " compression_depth, title, chunking_version, "
            " canonicalization_version, schema_version, ingest_ts) "
            "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
            ("c" * 64, "u1", "html", "surface", 0, "t",
             "tok-512-v1", "norm-v1", "v9.8.0", 0),
        )
        conn.execute(
            "INSERT INTO documents "
            "(document_root, document_uri, source_type, kind, "
            " compression_depth, title, chunking_version, "
            " canonicalization_version, schema_version, ingest_ts) "
            "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
            ("d" * 64, "u2", "html", "core", 0, "t",
             "tok-512-v1", "norm-v1", "v9.8.0", 0),
        )
        # Distillation process — should NOT match warrant-resolver-v1.
        conn.execute(
            "INSERT INTO derivations "
            "(core_root, src_root, process_id, proof_blob, distilled_at) "
            "VALUES (?, ?, ?, ?, ?)",
            ("d" * 64, "c" * 64, "tfidf-keywords-v1", b"{}", 0),
        )
        conn.commit()
    finally:
        conn.close()

    roots = warrant_chain_lookup(shards_dir)
    assert roots == frozenset()


def test_warrant_chain_lookup_empty_when_no_shards(tmp_path):
    """Nonexistent shards_dir → empty frozenset, no exception."""
    assert warrant_chain_lookup(tmp_path / "nope") == frozenset()


def test_warrant_chain_lookup_handles_missing_derivations_table(tmp_path):
    """Shards without a derivations table contribute zero rows
    silently (won't raise).
    """
    shards_dir = tmp_path / "shards"
    shards_dir.mkdir()
    bad_db = shards_dir / "000.db"
    conn = sqlite3.connect(str(bad_db))
    try:
        conn.execute("CREATE TABLE foo (x TEXT)")
        conn.commit()
    finally:
        conn.close()
    assert warrant_chain_lookup(shards_dir) == frozenset()


def test_warrant_chain_lookup_matches_alias_variant(tmp_path):
    """warrant-resolver-v1+alias is also a warrant-resolver match."""
    shards_dir = tmp_path / "shards"
    shards_dir.mkdir()
    db = shards_dir / "000.db"
    conn = sqlite3.connect(str(db))
    try:
        conn.executescript(SCHEMA_SQL)
        for r in ("e" * 64, "f" * 64):
            conn.execute(
                "INSERT INTO documents "
                "(document_root, document_uri, source_type, kind, "
                " compression_depth, title, chunking_version, "
                " canonicalization_version, schema_version, ingest_ts) "
                "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
                (r, "u_" + r[:4], "claim_pack", "surface", 0, "t",
                 "tok-512-v1", "norm-v1", "v9.8.0", 0),
            )
        conn.execute(
            "INSERT INTO derivations "
            "(core_root, src_root, process_id, proof_blob, distilled_at) "
            "VALUES (?, ?, ?, ?, ?)",
            ("e" * 64, "f" * 64, "warrant-resolver-v1+alias", b"{}", 0),
        )
        conn.commit()
    finally:
        conn.close()

    roots = warrant_chain_lookup(shards_dir)
    assert "e" * 64 in roots


# --- has_warrant_chain ----------------------------------------


def test_has_warrant_chain_basic():
    chain_set = frozenset({"a" * 64, "b" * 64})
    assert has_warrant_chain(["a" * 64, "z" * 64], chain_set) is True
    assert has_warrant_chain(["z" * 64], chain_set) is False
    assert has_warrant_chain([], chain_set) is False


def test_has_warrant_chain_empty_set_short_circuits():
    """Empty warrant_chain_roots → fast-fail; backward-compatible
    for callers that pass `frozenset()` as default.
    """
    assert has_warrant_chain(["a" * 64], frozenset()) is False


# --- verifier integration -------------------------------------


def _stub_chunk(source_root: str, span: str) -> dict:
    """Minimal chunk dict for build_evidence_map."""
    return {
        "source_root": source_root,
        "document_uri": "test://",
        "title": "t",
        "chunk_idx": 0,
        "chunk_root": "c" * 64,
        "span": span,
        "source_role": "primary_answer_source",
    }


def test_verify_claim_lattice_suppresses_warrant_missing_with_chain():
    """When the cited evidence's source_root is in
    warrant_chain_roots AND the lexical warrant_check would fail,
    the verifier suppresses WARRANT_MISSING and records the claim
    on warrant_proven_claim_idxs instead.

    The smoke regime: warrant_check fires WARRANT_MISSING only on
    relation-shaped questions ('who is X's boss?', 'when did Y
    happen?') — for those, an answer entity has to appear in the
    cited span. We construct exactly that case.
    """
    from arborist.qa.evidence import build_evidence_map
    from arborist.qa.verify import verify_claim_lattice

    chain_root = "a" * 64
    no_chain_root = "z" * 64

    # Evidence span has no proper-noun answer entity — so the
    # warrant_check will fire on a relation question.
    evidence_map = build_evidence_map(
        [_stub_chunk(chain_root, "some prose about a topic.")]
    )

    # Claim asserts a proper-noun ("Mr. Burns") that is NOT in the
    # cited span — exact failure shape that triggers WARRANT_MISSING.
    answer = "Mr. Burns is the boss. [E1]"

    # WITHOUT chain → WARRANT_MISSING fires.
    verdict_no_chain = verify_claim_lattice(
        answer,
        evidence_map,
        question="who is the boss?",
        warrant_chain_roots=frozenset(),
    )
    kinds_no = {v.get("kind") for v in verdict_no_chain.get("violations") or []}

    # WITH chain → WARRANT_MISSING suppressed; warrant_proven_claim_idxs
    # picks up the claim.
    verdict_with_chain = verify_claim_lattice(
        answer,
        evidence_map,
        question="who is the boss?",
        warrant_chain_roots=frozenset({chain_root}),
    )
    kinds_with = {v.get("kind") for v in verdict_with_chain.get("violations") or []}

    # The chain-set version should not produce WARRANT_MISSING for
    # the chain-rooted claim. (Other violations may still fire from
    # other rules — we only check that WARRANT_MISSING is suppressed
    # AND that warrant_proven_claim_idxs is populated.)
    if "WARRANT_MISSING" in kinds_no:
        # Fired without chain → chain-version should have suppressed.
        assert "WARRANT_MISSING" not in kinds_with or (
            verdict_with_chain.get("warrant_proven_claim_idxs") or []
        )
        # Chain side records the proven idx even if other rules fire.
        # Note: we don't strictly assert the absent-side fires —
        # warrant_check is shape-gated and may not trigger in this
        # narrow stub. We assert the chain-side behavior is correct
        # WHEN the absent-side fires.
    # Chain root NOT in set → no warrant_proven_claim_idxs.
    verdict_unrelated = verify_claim_lattice(
        answer,
        evidence_map,
        question="who is the boss?",
        warrant_chain_roots=frozenset({no_chain_root}),
    )
    assert (verdict_unrelated.get("warrant_proven_claim_idxs") or []) == []


def test_verify_returns_warrant_proven_idxs_field_always():
    """`warrant_proven_claim_idxs` is on every verdict (empty list
    if no chain-side suppression fired). Schema-level guarantee for
    downstream consumers."""
    from arborist.qa.evidence import build_evidence_map
    from arborist.qa.verify import (
        verify_claim_lattice,
        verify_claim_lattice_json,
    )

    evidence_map = build_evidence_map(
        [_stub_chunk("a" * 64, "some prose.")]
    )

    pointer_verdict = verify_claim_lattice(
        "X is Y. [E1]",
        evidence_map,
        question="what is X?",
    )
    assert "warrant_proven_claim_idxs" in pointer_verdict
    assert isinstance(pointer_verdict["warrant_proven_claim_idxs"], list)

    json_verdict = verify_claim_lattice_json(
        '{"claims":[{"text":"X is Y","evidence_ids":["E1"]}]}',
        evidence_map,
        question="what is X?",
    )
    assert "warrant_proven_claim_idxs" in json_verdict
    assert isinstance(json_verdict["warrant_proven_claim_idxs"], list)