PIG-Demo/envs/alert_triage/tests/test_triage_generator.py

"""What the generator promises about every world, checked over every digested seed."""

from __future__ import annotations

import json
from collections import Counter

import pytest

from alert_triage.generator import (
    BENIGN_PCT,
    RULE_IDS,
    VISIBLE_PCT,
    canonical_json,
    generate,
    is_held_out,
    screen_of,
)
from alert_triage.rng import days_from_civil

ALL = range(4096)


@pytest.fixture(scope="module")
def worlds():
    return [generate(s) for s in ALL]


def _day(iso: str) -> int:
    y, m, d = (int(x) for x in iso.split("-"))
    return days_from_civil(y, m, d)


def test_deterministic() -> None:
    assert canonical_json(generate(17)) == canonical_json(generate(17))
    assert canonical_json(generate(17)) != canonical_json(generate(18))


def test_canonical_json_has_no_floats(worlds) -> None:
    def walk(v):
        if isinstance(v, float):
            raise AssertionError("float in world")
        if isinstance(v, dict):
            for x in v.values():
                walk(x)
        if isinstance(v, list):
            for x in v:
                walk(x)
    for w in worlds[:512]:
        walk(w)
        assert json.loads(canonical_json(w)) == w


def test_mix_by_seed_is_close_to_declared(worlds) -> None:
    tiers = Counter(w["tier"] for w in worlds)
    n = len(worlds)
    assert abs(tiers["benign"] / n - BENIGN_PCT / 100) < 0.03
    assert abs(tiers["visible"] / n - VISIBLE_PCT / 100) < 0.03
    assert abs(tiers["hidden"] / n - (100 - BENIGN_PCT - VISIBLE_PCT) / 100) < 0.03


def test_every_rule_carries_both_labels(worlds) -> None:
    """CHEAT 3: no rule id may be a label. Each fires on benign AND suspicious seeds."""
    seen = Counter((w["alert"]["rule"], w["label"]) for w in worlds)
    for table in RULE_IDS.values():
        for rid in table.values():
            assert seen[(rid, "benign")] > 0, rid
            assert seen[(rid, "suspicious")] > 0, rid


def test_every_hidden_signal_lands_on_several_templates(worlds) -> None:
    """CHEAT 4: no fixed template -> lookup mapping to memorise."""
    pairs = Counter((w["overlay"], w["template"]) for w in worlds if w["overlay"])
    by_kind: dict[str, set[str]] = {}
    for (kind, template), _ in pairs.items():
        by_kind.setdefault(kind, set()).add(template)
    assert len(by_kind["funnel"]) == 6
    assert len(by_kind["serial_closer"]) >= 4
    assert len(by_kind["doc_mismatch"]) == 3


def test_overlay_only_on_benign_screens(worlds) -> None:
    for w in worlds:
        if w["overlay"]:
            assert w["template"].startswith("B") and w["tier"] == "hidden" and w["label"] == "suspicious"
        if w["template"].startswith("S"):
            assert w["label"] == "suspicious" and w["overlay"] is None


def test_nothing_is_dated_on_or_after_the_fire_date(worlds) -> None:
    for w in worlds:
        fired = _day(w["alert"]["fired"])
        assert all(_day(t["date"]) < fired for t in w["transactions"]), w["seed"]


def test_transaction_ids_are_chronological_and_dense(worlds) -> None:
    for w in worlds[:512]:
        ids = [t["id"] for t in w["transactions"]]
        assert ids == [f"T-{i + 1}" for i in range(len(ids))]
        days = [_day(t["date"]) for t in w["transactions"]]
        assert days == sorted(days)


def test_triggering_set_is_non_empty_and_shown(worlds) -> None:
    for w in worlds:
        ids = {t["id"] for t in w["transactions"]}
        assert w["alert"]["triggering"], w["seed"]
        assert set(w["alert"]["triggering"]) <= ids


def test_planted_is_citable_after_the_right_lookups(worlds) -> None:
    """Every planted id exists in the world: a kyc field, a T-id, a CP-id, a doc or a prior alert."""
    for w in worlds:
        universe = {f"kyc.{k}" for k in w["customer"]} | {t["id"] for t in w["transactions"]}
        universe |= {c["id"] for c in w["counterparties"]} | {d["id"] for d in w["documents"]} | {p["id"] for p in w["prior_alerts"]}
        assert w["planted"], w["seed"]
        for alt in w["planted"]:
            assert alt and set(alt) <= universe, (w["seed"], alt)


def test_benign_screens_always_carry_an_outbound_counterparty(worlds) -> None:
    for w in worlds:
        if w["template"].startswith("B"):
            assert any(c["direction"] == "out" for c in w["screen_counterparties"]), w["seed"]
            assert len(w["screen_counterparties"]) == 2
        assert 1 <= len(w["screen_counterparties"]) <= 2


def test_prior_alert_count_on_screen_matches_the_lookup(worlds) -> None:
    for w in worlds:
        assert len(w["prior_alerts"]) == w["customer"]["prior_alerts"]


def test_held_out_slice_moves_the_threshold_and_the_screen(worlds) -> None:
    held = [w for w in worlds if w["held_out"]]
    main = [w for w in worlds if not w["held_out"]]
    assert 0.10 < len(held) / len(worlds) < 0.15
    for w in held:
        assert is_held_out(w["seed"])
        assert w["currency"] == "KRN" and w["threshold"] == 2_000_000 and w["screen_order"] == "kyc_first"
        assert w["alert"]["rule"] in RULE_IDS["held_out"].values()
        assert w["customer"]["country"] == "Karnland"
    for w in main:
        assert w["currency"] == "USD" and w["threshold"] == 1_000_000 and w["screen_order"] == "alert_first"
        assert w["alert"]["rule"] in RULE_IDS["main"].values()


def test_structuring_band_is_threshold_relative(worlds) -> None:
    for w in worlds:
        if w["template"] == "S1":
            lo, hi = w["threshold"] * 92 // 100, w["threshold"] - 1
            deposits = [t for t in w["transactions"] if t["id"] in w["planted"][0]]
            assert len(deposits) == 8
            assert all(lo <= t["amount"] <= hi for t in deposits)
            assert len({t["branch"] for t in deposits}) == 3


def test_screen_carries_no_hidden_field(worlds) -> None:
    hidden = {"label", "typology", "planted", "overlay", "tier", "template", "reference_minutes", "reference_policy",
              "prior_alerts_list", "documents", "counterparties", "transactions", "other_alerted_accounts", "party", "rationale"}
    for w in worlds[:256]:
        text = canonical_json(screen_of(w))
        for key in hidden:
            assert f'"{key}"' not in text, key