arborist/tests/test_session_migrations.py

"""Schema migration idempotency for tickets #000014/#000017/#000020.

Verifies the three new ``_migrate_*`` functions in
``arborist.store.connect()``:

- ``_migrate_selfmodel_tables`` (selfmodel_records + selfmodel_capability_claims)
- ``_migrate_capital_ledger`` (capital_ledger)
- ``_migrate_memory_root`` (memory_records + memory_branch_summaries)

Each migration must be idempotent (re-running produces no errors) and
purely additive (existing tables/data unchanged).
"""

from __future__ import annotations

import hashlib
import json

import pytest

from arborist.store import (
    _migrate_capital_ledger,
    _migrate_memory_root,
    _migrate_selfmodel_tables,
    append_audit,
    connect,
    transaction,
)


_NEW_SESSION_TABLES = (
    "selfmodel_records",
    "selfmodel_capability_claims",
    "capital_ledger",
    "memory_records",
    "memory_branch_summaries",
)


def _list_tables(conn) -> set[str]:
    rows = conn.execute(
        "SELECT name FROM sqlite_master WHERE type='table'"
    ).fetchall()
    return {r["name"] for r in rows}


def test_fresh_db_has_all_session_tables(tmp_path):
    db = tmp_path / "shard.db"
    conn = connect(db)
    try:
        names = _list_tables(conn)
        for t in _NEW_SESSION_TABLES:
            assert t in names, f"missing table on fresh shard: {t}"
    finally:
        conn.close()


def test_re_connect_is_idempotent(tmp_path):
    """Opening the same shard twice must not raise."""
    db = tmp_path / "shard.db"
    conn1 = connect(db)
    conn1.close()
    conn2 = connect(db)
    try:
        names = _list_tables(conn2)
        for t in _NEW_SESSION_TABLES:
            assert t in names
    finally:
        conn2.close()


def test_explicit_migration_reapply_no_error(tmp_path):
    """Calling each ``_migrate_*`` helper twice on the same conn is a no-op."""
    db = tmp_path / "shard.db"
    conn = connect(db)
    try:
        _migrate_selfmodel_tables(conn)
        _migrate_selfmodel_tables(conn)
        _migrate_capital_ledger(conn)
        _migrate_capital_ledger(conn)
        _migrate_memory_root(conn)
        _migrate_memory_root(conn)
        # No exception → idempotency confirmed.
        names = _list_tables(conn)
        for t in _NEW_SESSION_TABLES:
            assert t in names
    finally:
        conn.close()


def test_session_tables_have_expected_schema(tmp_path):
    db = tmp_path / "shard.db"
    conn = connect(db)
    try:
        # Use PRAGMA table_info to check critical columns.
        sm_cols = {r["name"] for r in conn.execute(
            "PRAGMA table_info(selfmodel_records)"
        ).fetchall()}
        for col in (
            "selfmodel_root", "schema_version", "model_profile_hash",
            "verifier_method_root", "governance_policy_hash",
            "memory_root", "state", "body_blob", "audit_event_hash",
        ):
            assert col in sm_cols, f"selfmodel_records missing col {col}"

        cap_cols = {r["name"] for r in conn.execute(
            "PRAGMA table_info(capital_ledger)"
        ).fetchall()}
        for col in (
            "ledger_id", "audit_event_hash", "op_type",
            "living", "material", "financial", "intellectual",
            "experiential", "social", "cultural", "spiritual",
            "estimator_version", "estimator_inputs_blob", "recorded_at",
        ):
            assert col in cap_cols, f"capital_ledger missing col {col}"

        mem_cols = {r["name"] for r in conn.execute(
            "PRAGMA table_info(memory_records)"
        ).fetchall()}
        for col in (
            "memory_root", "schema_version", "audit_events_high_water",
            "branch_summaries_blob", "state", "audit_event_hash",
        ):
            assert col in mem_cols, f"memory_records missing col {col}"
    finally:
        conn.close()


def test_check_constraints_enforce_state_values(tmp_path):
    """``state`` columns reject anything outside live/stale/falsified."""
    import sqlite3

    db = tmp_path / "shard.db"
    conn = connect(db)
    try:
        with pytest.raises(sqlite3.IntegrityError):
            conn.execute(
                "INSERT INTO selfmodel_records ("
                " selfmodel_root, schema_version, model_profile_hash,"
                " verifier_method_root, governance_policy_hash,"
                " canonicalization_version, chunking_version,"
                " state, body_blob, audit_event_hash, created_at"
                ") VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
                (
                    "deadbeef" * 8, "selfmodel-v1", "mph", "vmr", "gph",
                    "norm-v1", "tok-512-v1",
                    "INVALID_STATE",  # rejected by CHECK constraint
                    b"{}", "ev", 0,
                ),
            )
    finally:
        conn.close()


def test_audit_chain_intact_after_session_writes(tmp_path):
    """All three new modules append events via ``store.append_audit``;
    the chain must re-verify byte-for-byte."""
    from arborist.capital import CapitalProfile, record as capital_record
    from arborist.memory import snapshot as memory_snapshot, store_snapshot as memory_store
    from arborist.selfmodel import snapshot as sm_snapshot, store_snapshot as sm_store

    db = tmp_path / "shard.db"
    conn = connect(db)
    try:
        with transaction(conn):
            # SelfModel snapshot.
            sm = sm_snapshot(conn)
            sm_store(conn, sm)
            # Memory snapshot.
            ms = memory_snapshot(conn)
            memory_store(conn, ms)
            # Capital ledger row tied to a fresh audit event.
            ev = append_audit(
                conn, event_type="ingest", subject_root=None,
                body={"docs": 1},
            )
            capital_record(
                conn, audit_event_hash=ev, op_type="ingest",
                profile=CapitalProfile(material=0.5),
            )

        rows = conn.execute(
            "SELECT event_hash, prev_event_hash, body FROM audit_events ORDER BY seq"
        ).fetchall()
        prev = None
        for row in rows:
            h = hashlib.sha256()
            if prev is not None:
                h.update(bytes.fromhex(prev))
            h.update(row["body"].encode("utf-8", errors="surrogatepass"))
            assert h.hexdigest() == row["event_hash"]
            assert row["prev_event_hash"] == prev
            prev = row["event_hash"]
    finally:
        conn.close()


def test_capital_writes_do_not_chain_into_audit_events(tmp_path):
    """Sibling-table invariant: capital_ledger rows must NOT participate
    in the audit chain. Recompute the chain ignoring capital rows;
    chain still verifies."""
    import sqlite3
    from arborist.capital import CapitalProfile, record as capital_record

    db = tmp_path / "shard.db"
    conn = connect(db)
    try:
        with transaction(conn):
            ev1 = append_audit(conn, event_type="ingest",
                               subject_root=None, body={"x": 1})
            capital_record(conn, audit_event_hash=ev1, op_type="ingest",
                           profile=CapitalProfile(material=1.0))
            ev2 = append_audit(conn, event_type="qa",
                               subject_root=None, body={"y": 2})
            # Multiple capital rows pointing at different events.
            capital_record(conn, audit_event_hash=ev2, op_type="qa",
                           profile=CapitalProfile(financial=0.05))
            capital_record(conn, audit_event_hash=ev2, op_type="qa",
                           profile=CapitalProfile(experiential=0.1))

        # capital_ledger has 3 rows.
        n_capital = conn.execute(
            "SELECT COUNT(*) FROM capital_ledger"
        ).fetchone()[0]
        assert n_capital == 3
        # audit_events has 2 rows.
        n_audit = conn.execute(
            "SELECT COUNT(*) FROM audit_events"
        ).fetchone()[0]
        assert n_audit == 2
    finally:
        conn.close()