arborist/tests/test_migration_audit_mode.py

"""Forward migration: legacy providence_cache shards add audit_mode columns.

Regression for a defect caught on 2026-04-28: SCHEMA_SQL referenced the
new audit_mode column inside a CREATE INDEX statement. CREATE TABLE IF
NOT EXISTS skips the table on legacy DBs, so the column doesn't yet exist
when the index statement runs — the whole executescript aborts.

The fix moved CREATE INDEX idx_providence_audit into _migrate_audit_mode.
This test asserts: a DB with the pre-audit-mode schema can be opened
through connect() & emerges with all four new columns + the index.
"""

from __future__ import annotations

import sqlite3

from arborist.store import connect


# Pre-v9.8-audit-mode providence_cache (snapshot of the schema before the
# audit_mode rollout). Mirrors the columns observed in real legacy shards.
LEGACY_SCHEMA = """
PRAGMA journal_mode = WAL;
CREATE TABLE providence_cache (
    cache_key                TEXT PRIMARY KEY,
    source_root              TEXT NOT NULL,
    document_uri             TEXT NOT NULL,
    question_hash            TEXT NOT NULL,
    question_text            TEXT NOT NULL,
    answer_text              TEXT NOT NULL,
    merkle_proof             TEXT NOT NULL,
    model_profile_hash       TEXT NOT NULL,
    conversation_hash        TEXT NOT NULL,
    governance_policy_hash   TEXT NOT NULL,
    schema_version           TEXT NOT NULL,
    canonicalization_version TEXT NOT NULL,
    chunking_version         TEXT NOT NULL,
    falsification_state      TEXT NOT NULL DEFAULT 'live',
    chain                    TEXT NOT NULL DEFAULT 'private',
    audit_event_hash         TEXT,
    created_at               INTEGER NOT NULL,
    last_hit_at              INTEGER,
    hit_count                INTEGER NOT NULL DEFAULT 0
);
"""


def _seed_legacy(path) -> None:
    raw = sqlite3.connect(path)
    raw.executescript(LEGACY_SCHEMA)
    # One pre-existing row simulating a real legacy record.
    raw.execute(
        "INSERT INTO providence_cache "
        "(cache_key, source_root, document_uri, question_hash, question_text, "
        " answer_text, merkle_proof, model_profile_hash, conversation_hash, "
        " governance_policy_hash, schema_version, canonicalization_version, "
        " chunking_version, audit_event_hash, created_at) "
        "VALUES ('legacy_key', 'root', 'corpus://x', 'qh', 'Q?', 'A.', '{}', "
        " 'mh', 'ch', 'gh', 'v9.8.0', 'norm-v1', 'tok-512-v1', NULL, 0)"
    )
    raw.commit()
    raw.close()


def test_legacy_providence_cache_migrates(tmp_path):
    db = tmp_path / "legacy.db"
    _seed_legacy(db)

    conn = connect(db)
    try:
        cols = {r["name"] for r in conn.execute("PRAGMA table_info(providence_cache)")}
        assert {"audit_mode", "n_quotes", "n_verified", "unverified_quotes"} <= cols

        # Pre-existing row inherits the safest default — UNGROUNDED.
        row = conn.execute(
            "SELECT audit_mode, n_quotes, n_verified, unverified_quotes "
            "FROM providence_cache WHERE cache_key='legacy_key'"
        ).fetchone()
        assert row["audit_mode"] == "UNGROUNDED"
        assert row["n_quotes"] == 0
        assert row["n_verified"] == 0
        assert row["unverified_quotes"] is None

        # Index is in place after migration.
        idx_names = {
            r["name"]
            for r in conn.execute(
                "SELECT name FROM sqlite_master WHERE type='index'"
            )
        }
        assert "idx_providence_audit" in idx_names
    finally:
        conn.close()


def test_migration_idempotent(tmp_path):
    """Two consecutive opens of the same legacy DB don't ALTER TABLE twice."""
    db = tmp_path / "twice.db"
    _seed_legacy(db)

    connect(db).close()
    # Second open: PRAGMA table_info already shows the columns; ALTER TABLE
    # would error with "duplicate column name" if we didn't gate it.
    conn = connect(db)
    try:
        cols = {r["name"] for r in conn.execute("PRAGMA table_info(providence_cache)")}
        assert "audit_mode" in cols
    finally:
        conn.close()