stack/tests/lead/test_tools.py

"""The agent writes ledger rows. Databricks stays behind an injected transport."""

from __future__ import annotations

from contextlib import contextmanager

import pytest

from lead.ledger.store import append, open_ledger
from lead.tools import (
    failures_check,
    features_propose,
    features_validate,
    findings_record,
    ledger_get_claim,
    ledger_search,
    limits_compute,
    review_trigger_search,
    runs_compare,
    runs_get,
    runs_start,
    targets_list,
    tracer_factory,
)
from perf import tracer


def _claim(con, claim_id: str = "claim-1", subject: str = "Frailty index") -> None:
    append(
        con,
        "claim",
        {
            "id": claim_id,
            "kind": "precursor",
            "subject": subject,
            "target_id": "Q.ACR",
            "tiers": "literature",
            "content_hash": claim_id,
            "actor": "tester",
        },
    )


def test_the_default_tracer_factory_is_perf():
    assert tracer_factory is tracer


def test_ledger_search_filters_and_is_traced(tmp_path, monkeypatch):
    class Recording:
        def __init__(self) -> None:
            self.names: list[str] = []

        def __call__(self, scope: str) -> Recording:
            assert scope == "lead.tools"
            return self

        @contextmanager
        def start_as_current_span(self, name: str, **_kwargs: object):
            self.names.append(name)
            yield None

    recording = Recording()
    monkeypatch.setattr("lead.tools.tracer_factory", recording)
    con = open_ledger(str(tmp_path / "ledger.duckdb"))
    try:
        _claim(con, "claim-1", "Frailty index")
        _claim(con, "claim-2", "Other topic")
        found = ledger_search(con, target_id="Q.ACR", text="FRAIL")
        assert [row["id"] for row in found] == ["claim-1"]
        assert ledger_get_claim(con, "claim-1")["subject"] == "Frailty index"
        assert ledger_get_claim(con, "missing") is None
    finally:
        con.close()
    assert "ledger_search" in recording.names
    assert "ledger_get_claim" in recording.names


def test_targets_include_quality_and_high_needs():
    ids = {target.id for target in targets_list()}
    assert "Q.ACR" in ids
    assert "HN.4" in ids


def test_limits_compute_dispatches_the_three_expressions():
    assert limits_compute("headroom", prior=1.0, caught=1.05).scored == 1.04
    assert limits_compute("normalization", benchmark=1000.0, factor=1.25) == 800.0
    assert limits_compute("shrinkage", aco_n=100, prior_n=300) == 0.25
    with pytest.raises(ValueError):
        limits_compute("nope")


def test_a_feature_needs_a_known_target_and_a_stored_claim(tmp_path):
    con = open_ledger(str(tmp_path / "ledger.duckdb"))
    try:
        with pytest.raises(KeyError):
            features_validate(con, target_id="NO.SUCH", claim_ids=("claim-1",))
        with pytest.raises(ValueError):
            features_validate(con, target_id="Q.ACR", claim_ids=())
        with pytest.raises(ValueError):
            features_validate(con, target_id="Q.ACR", claim_ids=("missing",))
        _claim(con)
        assert (
            features_propose(
                con,
                target_id="Q.ACR",
                claim_ids=("claim-1",),
                expression="kim_index",
                actor="tester",
            )
            is True
        )
        assert (
            features_propose(
                con,
                target_id="Q.ACR",
                claim_ids=("claim-1",),
                expression="kim_index",
                actor="other",
            )
            is False
        )
    finally:
        con.close()


def test_local_runs_do_not_call_a_transport(tmp_path):
    con = open_ledger(str(tmp_path / "ledger.duckdb"))

    def boom(_row: dict) -> None:
        raise AssertionError("transport called")

    try:
        run_id = runs_start(
            con,
            where_ran="local",
            target_id="Q.ACR",
            actor="tester",
            model="glm",
            transport=boom,
            finished_on="2026-10-31",
            metrics="mae=1",
        )
        assert runs_get(con, run_id)["metrics"] == "mae=1"
    finally:
        con.close()


def test_databricks_refuses_before_calling_a_transport(tmp_path):
    con = open_ledger(str(tmp_path / "ledger.duckdb"))
    called: list[dict] = []
    try:
        with pytest.raises(RuntimeError):
            runs_start(
                con,
                where_ran="databricks",
                target_id="Q.ACR",
                actor="tester",
                model="glm",
                transport=None,
                spent=0,
                budget=1,
            )
        with pytest.raises(RuntimeError):
            runs_start(
                con,
                where_ran="databricks",
                target_id="Q.ACR",
                actor="tester",
                model="glm",
                transport=called.append,
                spent=1,
                budget=1,
            )
        assert called == []
        assert con.execute("SELECT count(*) FROM lead_evidence.run").fetchone()[0] == 0
    finally:
        con.close()


def test_a_repeated_databricks_start_does_not_call_the_transport_again(tmp_path):
    con = open_ledger(str(tmp_path / "ledger.duckdb"))
    called: list[dict] = []
    kwargs = {
        "where_ran": "databricks",
        "target_id": "Q.ACR",
        "actor": "tester",
        "model": "glm",
        "finished_on": "2026-10-31",
        "metrics": "mae=1",
        "spent": 0,
        "budget": 1,
    }
    try:
        first = runs_start(con, transport=called.append, **kwargs)
        second = runs_start(con, transport=called.append, **kwargs)
        assert first == second
        assert len(called) == 1
    finally:
        con.close()


def test_an_injected_databricks_transport_records_the_run(tmp_path):
    con = open_ledger(str(tmp_path / "ledger.duckdb"))
    called: list[dict] = []
    try:
        left = runs_start(
            con,
            where_ran="databricks",
            target_id="Q.ACR",
            actor="tester",
            model="glm",
            transport=called.append,
            spent=0,
            budget=1,
            finished_on="2026-10-31",
            metrics="mae=1",
        )
        right = runs_start(
            con,
            where_ran="local",
            target_id="Q.ACR",
            actor="tester",
            model="glm",
            finished_on="2026-11-01",
            metrics="mae=2",
            gate_results="baseline",
        )
        assert len(called) == 1
        assert runs_compare(con, left, right)["left_metrics"] == "mae=1"
        assert runs_compare(con, left, right)["right_metrics"] == "mae=2"
        failed = failures_check(con, "Q.ACR")
        assert [row["id"] for row in failed] == [right]
        with pytest.raises(KeyError):
            runs_compare(con, left, "missing")
    finally:
        con.close()


def test_findings_and_search_triggers_write_only_the_ledger(tmp_path):
    con = open_ledger(str(tmp_path / "ledger.duckdb"))
    try:
        assert (
            findings_record(
                con,
                {
                    "relation": "contradicts",
                    "numbers": "1",
                    "content_hash": "finding-1",
                    "actor": "tester",
                },
            )
            is True
        )
        assert review_trigger_search(
            con, project="lead-precursors", now="2026-10-31"
        ) == ("nightly",)
    finally:
        con.close()