tests/test_predictor_per_layer_cap.py

"""Tests for CascadePredictor._apply_per_layer_cap — Step C issue #11.

Two layers of test:

1. Pure unit tests of the helper itself (no LLM, no retrieval).
2. Trace-record contract tests: confirm the predictor surfaces
   ``truncated_by_per_layer_cap`` per layer in the BFS trace.

Integration via real LLM is out of scope here — that's covered by the
end-to-end gold evaluation in scripts/05_evaluate.py.
"""

from __future__ import annotations

from src.models.schemas import CascadeNode
from src.rag.predictor import CascadePredictor


def _node(nid: str) -> CascadeNode:
    return CascadeNode(
        id=nid,
        description=f"node {nid}",
        domain="infrastructure/power",
        severity="high",
        time_offset_hours=10.0,
        mechanism="m",
        parent_ids=[],
    )


class _CapHarness(CascadePredictor):
    """Bare CascadePredictor that skips the heavy ChromaDB / LLM init.

    Bypasses ``__init__`` and only sets ``max_new_nodes_per_layer`` and
    ``dump_full_trace`` so we can unit-test the cap + trace-shape helpers
    in isolation.
    """

    def __init__(self, cap: int = 10):
        # NOTE: deliberately do NOT call super().__init__ — that would
        # require a working LLM client + ChromaDB + index. We're only
        # exercising the pure helper here.
        self.max_new_nodes_per_layer = cap
        self.dump_full_trace = False


# ---------------------------------------------------------------------------
# 1. Pure helper tests
# ---------------------------------------------------------------------------
def test_under_cap_returns_unchanged():
    p = _CapHarness(cap=10)
    nodes = [_node(f"E{i}") for i in range(5)]
    kept, dropped = p._apply_per_layer_cap(nodes)
    assert kept == nodes
    assert dropped == 0


def test_at_cap_returns_unchanged():
    p = _CapHarness(cap=5)
    nodes = [_node(f"E{i}") for i in range(5)]
    kept, dropped = p._apply_per_layer_cap(nodes)
    assert len(kept) == 5
    assert dropped == 0


def test_over_cap_truncates_by_emit_order_when_no_confidence():
    p = _CapHarness(cap=3)
    nodes = [_node(f"E{i}") for i in range(7)]
    kept, dropped = p._apply_per_layer_cap(nodes)
    assert dropped == 4
    assert [n.id for n in kept] == ["E0", "E1", "E2"]


def test_over_cap_truncates_by_confidence_descending():
    p = _CapHarness(cap=3)
    nodes = [_node(f"E{i}") for i in range(7)]
    confidence = {
        "E0": 0.10, "E1": 0.95, "E2": 0.20,
        "E3": 0.99, "E4": 0.30, "E5": 0.85, "E6": 0.05,
    }
    kept, dropped = p._apply_per_layer_cap(nodes, confidence)
    assert dropped == 4
    kept_ids = {n.id for n in kept}
    assert kept_ids == {"E1", "E3", "E5"}  # top-3 confidence


def test_truncation_preserves_original_emit_order_for_kept_set():
    """Stable display: kept nodes appear in the order LLM emitted them,
    not the confidence-ranked order. Trace + UI display stays predictable."""
    p = _CapHarness(cap=3)
    nodes = [_node(f"E{i}") for i in range(5)]
    confidence = {"E0": 0.5, "E1": 0.99, "E2": 0.6, "E3": 0.8, "E4": 0.7}
    kept, _dropped = p._apply_per_layer_cap(nodes, confidence)
    # top-3 by confidence: E1 (0.99), E3 (0.8), E4 (0.7)
    # original order should be preserved → E1, E3, E4 (not E1, E4, E3)
    assert [n.id for n in kept] == ["E1", "E3", "E4"]


def test_empty_confidence_falls_back_to_emit_order():
    p = _CapHarness(cap=2)
    nodes = [_node(f"E{i}") for i in range(5)]
    kept, dropped = p._apply_per_layer_cap(nodes, {})
    assert dropped == 3
    assert [n.id for n in kept] == ["E0", "E1"]


def test_partial_confidence_treats_missing_as_zero():
    """Nodes without confidence entries get 0.0 — they sort last."""
    p = _CapHarness(cap=2)
    nodes = [_node(f"E{i}") for i in range(5)]
    # Only E2 and E4 have explicit confidence
    confidence = {"E2": 0.9, "E4": 0.7}
    kept, _dropped = p._apply_per_layer_cap(nodes, confidence)
    kept_ids = {n.id for n in kept}
    assert kept_ids == {"E2", "E4"}


# ---------------------------------------------------------------------------
# 2. _record_layer carries the truncation count
# ---------------------------------------------------------------------------
def test_record_layer_default_truncation_zero():
    p = _CapHarness()
    rec = p._record_layer(
        layer_idx=0,
        frontier_ids=[],
        evidence={},
        stop_reason=None,
        produced_ids=["E1"],
    )
    assert rec["truncated_by_per_layer_cap"] == 0


def test_record_layer_truncation_count_propagates():
    p = _CapHarness()
    rec = p._record_layer(
        layer_idx=2,
        frontier_ids=["E1", "E2"],
        evidence={},
        stop_reason=None,
        produced_ids=["E3", "E4", "E5"],
        truncated_by_per_layer_cap=7,
    )
    assert rec["truncated_by_per_layer_cap"] == 7
    assert rec["produced_ids"] == ["E3", "E4", "E5"]
    assert rec["layer"] == 2