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

"""Rejections, charges, shown ids, and the abort/truncate line."""

from __future__ import annotations

import json

from alert_triage.engine import Engine
from alert_triage.generator import MAX_TURNS, world_for_seed
from alert_triage.reward import DOC_MINUTES, LOOKUP_MINUTES, SCREEN_MINUTES, TURN_MINUTES, score

JUNK = ["", "   ", None, "no json here", "{", "}", "{}", "[]", "null", "{\"action\":null}", "{\"action\":\"lookup\"}",
        "{\"action\":\"lookup\",\"what\":\"history\"}", "{\"action\":\"lookup\",\"what\":\"history\",\"month\":7}",
        "{\"action\":\"lookup\",\"what\":\"counterparty\"}", "{\"action\":\"lookup\",\"what\":\"counterparty\",\"id\":[]}",
        "{\"action\":\"close\"}", "{\"action\":\"close\",\"reason\":\"RULE_ARTEFACT\"}",
        "{\"action\":\"close\",\"reason\":\"RULE_ARTEFACT\",\"cites\":{}}", "{\"action\":\"escalate\",\"cites\":[\"kyc.country\"]}",
        "{\"action\":\"lookup\",\"what\":\"documents\",\"x\":NaN}", "\x00\x01", "{" * 50]


def _e(seed: int = 0) -> Engine:
    return Engine(world_for_seed(seed))


def test_screen_costs_and_first_turn() -> None:
    e = _e()
    assert e.minutes == SCREEN_MINUTES and e.turns == 0 and not e.done


def test_junk_never_raises_and_costs_one_turn_each() -> None:
    for junk in JUNK:
        e = _e()
        step = e.step(junk)
        assert not step.accepted and step.reason and e.turns == 1 and e.rejected == 1
        assert e.minutes == SCREEN_MINUTES + TURN_MINUTES


def test_trailing_garbage_after_a_balanced_object_is_fine() -> None:
    e = _e()
    assert e.step('{"action":"lookup","what":"documents"}' + "}" * 50 + " and more prose").accepted


def test_lookup_charges_its_price_only_when_accepted() -> None:
    e = _e()
    e.step('{"action":"lookup","what":"documents"}')
    assert e.minutes == SCREEN_MINUTES + TURN_MINUTES + DOC_MINUTES
    e.step('{"action":"lookup","what":"documents"}')       # repeat: rejected
    assert e.minutes == SCREEN_MINUTES + 2 * TURN_MINUTES + DOC_MINUTES and e.rejected == 1
    e.step('{"action":"lookup","what":"prior_alerts"}')
    assert e.minutes == SCREEN_MINUTES + 3 * TURN_MINUTES + DOC_MINUTES + LOOKUP_MINUTES


def test_history_window_is_the_twelve_months_ending_in_the_fire_month() -> None:
    e = _e()
    months = e.world["months"]
    assert len(months) == 12 and months[-1] == e.world["alert"]["fired"][:7]
    assert e.step(json.dumps({"action": "lookup", "what": "history", "month": months[0]})).accepted
    before = months[0][:4] + "-" + f"{int(months[0][5:]) - 1:02d}" if months[0][5:] != "01" else f"{int(months[0][:4]) - 1}-12"
    assert not e.step(json.dumps({"action": "lookup", "what": "history", "month": before})).accepted


def test_counterparty_must_have_been_shown() -> None:
    e = _e()
    on_screen = e.world["screen_counterparties"][0]["id"]
    assert e.step(json.dumps({"action": "lookup", "what": "counterparty", "id": on_screen})).accepted
    assert not e.step('{"action":"lookup","what":"counterparty","id":"CP-999"}').accepted
    assert not e.step(json.dumps({"action": "lookup", "what": "counterparty", "id": on_screen})).accepted


def test_cites_must_be_shown_and_are_deduplicated() -> None:
    e = _e()
    step = e.step('{"action":"close","reason":"RULE_ARTEFACT","cites":["kyc.country","doc.D-0000"]}')
    assert not step.accepted and "not been shown" in step.reason
    step = e.step('{"action":"close","reason":"RULE_ARTEFACT","cites":["kyc.country","kyc.country","kyc.pep"]}')
    assert step.accepted and step.action["cites"] == ["kyc.country", "kyc.pep"] and e.done


def test_documents_become_citable_after_the_lookup() -> None:
    e = _e()
    doc = e.world["documents"][0]["id"]
    assert not e.step(json.dumps({"action": "close", "reason": "DOCUMENTED_SOURCE_OF_FUNDS", "cites": [doc]})).accepted
    assert e.step('{"action":"lookup","what":"documents"}').accepted
    assert e.step(json.dumps({"action": "close", "reason": "DOCUMENTED_SOURCE_OF_FUNDS", "cites": [doc]})).accepted


def test_enums_are_exact_case() -> None:
    e = _e()
    assert not e.step('{"action":"close","reason":"rule_artefact","cites":["kyc.country"]}').accepted
    assert not e.step('{"action":"ESCALATE","typology":"UNKNOWN","cites":["kyc.country"]}').accepted
    assert e.step('{"action":"escalate","typology":"UNKNOWN","cites":["kyc.country"]}').accepted


def test_eight_rejections_abort_and_score_zero() -> None:
    e = _e()
    for _ in range(MAX_TURNS):
        e.step("nothing")
    assert e.done and e.aborted and e.outcome == "aborted"
    assert e.step("{}").reason == "the episode is over" and e.turns == MAX_TURNS
    assert score(e.episode()) == {"caught": 0.0, "hours": 0.0, "evidence": 0.0}


def test_a_disposition_on_the_eighth_turn_is_accepted() -> None:
    e = _e()
    for _ in range(MAX_TURNS - 1):
        e.step("nothing")
    assert e.step('{"action":"close","reason":"RULE_ARTEFACT","cites":["kyc.country"]}').accepted
    assert e.outcome in ("solved", "failed")


def test_stopping_early_is_truncated_not_aborted() -> None:
    e = _e()
    e.step("nothing")
    e.step('{"action":"lookup","what":"documents"}')
    assert not e.done and e.outcome is None
    ep = e.episode()
    assert ep.truncated
    assert score(ep) == {"caught": None, "hours": None, "evidence": None}


def test_view_never_shows_hidden_fields() -> None:
    e = _e(5)
    e.step('{"action":"lookup","what":"documents"}')
    e.step('{"action":"lookup","what":"prior_alerts"}')
    text = json.dumps(e.view())
    for key in ("label", "planted", "typology", "overlay", "tier", "template", "reference_minutes", "reference_policy"):
        assert f'"{key}"' not in text, key