arborist/tests/test_prometheus_audit.py

"""Phase 2 advisory-write tests for ticket #000037.

Covers the sibling-table ``controller_events`` migration and the
``emit_controller_events`` helper in
``arborist.substrate.prometheus_audit``.

These tests use an inline stub ``ControllerDecision`` rather than
importing the real one from ``arborist.substrate.prometheus`` —
Phase 1 (the controller function) is being implemented in
parallel and may not be present when these tests run. The helper
only reads ControllerDecision attributes (it never executes
decision logic), so a duck-typed stub is sufficient.

Invariants asserted:
- Migration creates the ``controller_events`` table with the
  expected schema + indexes.
- Migration is idempotent across repeated ``connect()`` calls.
- ``emit_controller_events`` writes one decision row, one
  difficulty row, and N allocation rows (one per nonzero
  allocation branch).
- Re-emitting the same ControllerDecision is a no-op (no duplicate
  rows under the ``UNIQUE (event_kind, body_hash)`` constraint).
- The audit chain is NOT mutated by any controller_events write.
- ``canonical_body`` is deterministic and sha256-hex.
"""

from __future__ import annotations

import hashlib
import json
from dataclasses import dataclass

import pytest

from arborist.store import (
    append_audit,
    connect,
    invalidate_migration_cache,
    latest_event_hash,
    transaction,
)
from arborist.substrate.prometheus_audit import (
    canonical_body,
    emit_controller_events,
)


# --- stub ControllerDecision ------------------------------------------
#
# Duck-typed; matches the Phase 1 contract. Constructed inline per
# test so Phase 2 tests don't depend on Phase 1's implementation
# landing.


@dataclass(frozen=True)
class StubDecision:
    selected_branch_id: str | None
    label: str
    allocations: dict
    difficulty_next: float
    veto_reasons: dict
    entropy: float
    notes: tuple
    memory_proposals: tuple = ()
    selfmodel_proposals: tuple = ()
    advisory_events: tuple = ()


def _stub_decision(**overrides) -> StubDecision:
    """Build a stub ControllerDecision with sensible defaults."""
    base = dict(
        selected_branch_id="B0",
        label="ACCEPT",
        allocations={"B0": 0.7, "B1": 0.3, "B2": 0.0},
        difficulty_next=0.42,
        veto_reasons={"B2": ("low_score",)},
        entropy=1.23,
        notes=("note-a",),
    )
    base.update(overrides)
    return StubDecision(**base)


# --- fixtures ---------------------------------------------------------


@pytest.fixture
def shard(tmp_path):
    """Fresh shard.db per test. Migration runs on first connect."""
    db_path = tmp_path / "shard.db"
    invalidate_migration_cache(db_path)
    conn = connect(db_path)
    try:
        yield conn, db_path
    finally:
        conn.close()
        invalidate_migration_cache(db_path)


# --- migration --------------------------------------------------------


def test_migration_creates_controller_events_table(shard):
    conn, _ = shard
    row = conn.execute(
        "SELECT name FROM sqlite_master "
        "WHERE type='table' AND name='controller_events'"
    ).fetchone()
    assert row is not None

    cols = {r["name"] for r in conn.execute(
        "PRAGMA table_info(controller_events)"
    )}
    expected = {
        "event_id",
        "organism_root",
        "branch_id",
        "event_kind",
        "label",
        "entropy",
        "difficulty",
        "allocation",
        "body_blob",
        "body_hash",
        "recorded_at",
    }
    assert expected.issubset(cols)

    # Indexes present.
    idx = {r["name"] for r in conn.execute(
        "SELECT name FROM sqlite_master "
        "WHERE type='index' AND tbl_name='controller_events'"
    )}
    assert "idx_controller_events_organism" in idx
    assert "idx_controller_events_kind" in idx
    assert "idx_controller_events_at" in idx


def test_migration_idempotent(tmp_path):
    db_path = tmp_path / "shard.db"
    invalidate_migration_cache(db_path)
    conn1 = connect(db_path)
    conn1.close()
    invalidate_migration_cache(db_path)  # Force re-probe path on second open.
    conn2 = connect(db_path)
    try:
        row = conn2.execute(
            "SELECT name FROM sqlite_master "
            "WHERE type='table' AND name='controller_events'"
        ).fetchone()
        assert row is not None
        # No row count drift — table is empty either way.
        n = conn2.execute(
            "SELECT COUNT(*) FROM controller_events"
        ).fetchone()[0]
        assert n == 0
    finally:
        conn2.close()
        invalidate_migration_cache(db_path)


# --- emit_controller_events ------------------------------------------


def test_emit_writes_decision_row(shard):
    conn, _ = shard
    dec = _stub_decision()
    with transaction(conn):
        ids = emit_controller_events(conn, dec, "org-root-abc")
    assert len(ids) >= 1
    row = conn.execute(
        "SELECT * FROM controller_events "
        "WHERE event_kind='controller_decision'"
    ).fetchone()
    assert row is not None
    assert row["organism_root"] == "org-root-abc"
    assert row["branch_id"] == "B0"
    assert row["label"] == "ACCEPT"
    assert row["entropy"] == pytest.approx(1.23)
    assert row["body_hash"] is not None
    assert len(row["body_hash"]) == 64  # sha256 hex
    # Body parses back to a dict with the right shape.
    body = json.loads(row["body_blob"])
    assert body["kind"] == "controller_decision"
    assert body["selected_branch_id"] == "B0"
    assert body["label"] == "ACCEPT"


def test_emit_writes_difficulty_row(shard):
    conn, _ = shard
    dec = _stub_decision(difficulty_next=0.91)
    with transaction(conn):
        emit_controller_events(conn, dec, "org-1")
    row = conn.execute(
        "SELECT * FROM controller_events "
        "WHERE event_kind='controller_difficulty'"
    ).fetchone()
    assert row is not None
    assert row["organism_root"] == "org-1"
    assert row["branch_id"] is None
    assert row["difficulty"] == pytest.approx(0.91)
    body = json.loads(row["body_blob"])
    assert body["kind"] == "controller_difficulty"
    assert body["difficulty_next"] == pytest.approx(0.91)


def test_emit_writes_allocation_rows_for_nonzero_branches(shard):
    conn, _ = shard
    dec = _stub_decision(
        allocations={"B0": 0.5, "B1": 0.5, "B2": 0.0, "B3": 0.0}
    )
    with transaction(conn):
        emit_controller_events(conn, dec, "org-X")
    rows = conn.execute(
        "SELECT branch_id, allocation FROM controller_events "
        "WHERE event_kind='controller_budget_allocation' "
        "ORDER BY branch_id"
    ).fetchall()
    assert len(rows) == 2
    branches = {r["branch_id"] for r in rows}
    assert branches == {"B0", "B1"}
    for r in rows:
        assert r["allocation"] == pytest.approx(0.5)


def test_emit_idempotent_same_decision(shard):
    conn, _ = shard
    dec = _stub_decision()
    with transaction(conn):
        first = emit_controller_events(conn, dec, "org-r")
    with transaction(conn):
        second = emit_controller_events(conn, dec, "org-r")
    assert len(first) >= 3  # 1 decision + 1 difficulty + 2 nonzero allocs
    assert second == []  # all dupes — UNIQUE (event_kind, body_hash)
    # Row counts match the first emit exactly.
    n = conn.execute(
        "SELECT COUNT(*) FROM controller_events"
    ).fetchone()[0]
    assert n == len(first)


def test_emit_does_not_touch_audit_events_chain(shard):
    conn, _ = shard
    # Baseline: write one audit row, snapshot the head.
    with transaction(conn):
        baseline_hash = append_audit(
            conn,
            event_type="test_baseline",
            body={"k": "v"},
            subject_root="root-1",
        )
    head_before = latest_event_hash(conn)
    n_audit_before = conn.execute(
        "SELECT COUNT(*) FROM audit_events"
    ).fetchone()[0]
    assert head_before == baseline_hash

    # Emit controller events.
    dec = _stub_decision()
    with transaction(conn):
        emit_controller_events(conn, dec, "org-untouched")

    # Audit chain unchanged.
    head_after = latest_event_hash(conn)
    n_audit_after = conn.execute(
        "SELECT COUNT(*) FROM audit_events"
    ).fetchone()[0]
    assert head_after == head_before
    assert n_audit_after == n_audit_before


def test_canonical_body_deterministic():
    body = {"b": 2, "a": 1, "nested": {"y": "yes", "x": "no"}}
    s1, h1 = canonical_body(body)
    s2, h2 = canonical_body(dict(body))  # different dict identity, same data
    assert s1 == s2
    assert h1 == h2


def test_canonical_body_distinct_for_distinct_input():
    s1, h1 = canonical_body({"a": 1})
    s2, h2 = canonical_body({"a": 2})
    assert s1 != s2
    assert h1 != h2


def test_emit_handles_empty_allocations(shard):
    conn, _ = shard
    dec = _stub_decision(allocations={})
    with transaction(conn):
        ids = emit_controller_events(conn, dec, "org-empty")
    # Decision + difficulty rows only, no allocation rows.
    assert len(ids) == 2
    n_alloc = conn.execute(
        "SELECT COUNT(*) FROM controller_events "
        "WHERE event_kind='controller_budget_allocation'"
    ).fetchone()[0]
    assert n_alloc == 0


def test_emit_handles_none_selected_branch(shard):
    conn, _ = shard
    dec = _stub_decision(
        selected_branch_id=None,
        label="REJECT",
        allocations={},
        veto_reasons={"B0": ("low_score",), "B1": ("low_score",)},
    )
    with transaction(conn):
        emit_controller_events(conn, dec, "org-rej")
    row = conn.execute(
        "SELECT * FROM controller_events "
        "WHERE event_kind='controller_decision'"
    ).fetchone()
    assert row is not None
    assert row["branch_id"] is None
    assert row["label"] == "REJECT"


def test_body_hash_is_sha256_hex(shard):
    conn, _ = shard
    dec = _stub_decision()
    with transaction(conn):
        emit_controller_events(conn, dec, "org-h")
    rows = conn.execute(
        "SELECT body_blob, body_hash FROM controller_events"
    ).fetchall()
    assert rows
    for r in rows:
        blob = r["body_blob"]
        expected = hashlib.sha256(
            blob.encode("utf-8", errors="surrogatepass")
        ).hexdigest()
        assert r["body_hash"] == expected
        # hex characters only, length 64
        assert len(r["body_hash"]) == 64
        int(r["body_hash"], 16)  # parses as hex


def test_query_by_organism_root_returns_rows(shard):
    conn, _ = shard
    dec_a = _stub_decision(selected_branch_id="A", notes=("a-only",))
    dec_b = _stub_decision(selected_branch_id="B", notes=("b-only",))
    with transaction(conn):
        emit_controller_events(conn, dec_a, "org-alpha")
        emit_controller_events(conn, dec_b, "org-beta")
    alpha = conn.execute(
        "SELECT COUNT(*) FROM controller_events "
        "WHERE organism_root=?",
        ("org-alpha",),
    ).fetchone()[0]
    beta = conn.execute(
        "SELECT COUNT(*) FROM controller_events "
        "WHERE organism_root=?",
        ("org-beta",),
    ).fetchone()[0]
    assert alpha > 0
    assert beta > 0
    # Each shard wrote >= 1 decision row.
    n_alpha_dec = conn.execute(
        "SELECT COUNT(*) FROM controller_events "
        "WHERE organism_root=? AND event_kind='controller_decision'",
        ("org-alpha",),
    ).fetchone()[0]
    assert n_alpha_dec == 1


def test_query_by_kind_returns_rows(shard):
    conn, _ = shard
    dec = _stub_decision(
        allocations={"X": 0.6, "Y": 0.4},
    )
    with transaction(conn):
        emit_controller_events(conn, dec, "org-K")
    kinds = {
        r["event_kind"]
        for r in conn.execute(
            "SELECT DISTINCT event_kind FROM controller_events"
        ).fetchall()
    }
    assert "controller_decision" in kinds
    assert "controller_difficulty" in kinds
    assert "controller_budget_allocation" in kinds
    n_alloc = conn.execute(
        "SELECT COUNT(*) FROM controller_events "
        "WHERE event_kind='controller_budget_allocation'"
    ).fetchone()[0]
    assert n_alloc == 2