PTSRaiAgent/tests/test_safety.py

"""Safety detector tests (issue #4). Recall on acute phrases must be 100%."""

import time
from pathlib import Path

import pytest
import yaml

from app.core.privacy import scrub
from app.core.safety import (
    CrisisScript,
    SafetyResult,
    check,
    crisis_reply,
    load_lexicon,
    normalise,
)
from app.core.types import RiskLevel, ScrubbedText

CASES_PATH = Path(__file__).parent / "data" / "safety_cases.yaml"
CASES: list[dict] = yaml.safe_load(CASES_PATH.read_text(encoding="utf-8"))["cases"]
LABELS = {"acute": RiskLevel.ACUTE, "possible": RiskLevel.POSSIBLE, "none": RiskLevel.NONE}


def _ids(cases: list[dict]) -> list[str]:
    return [c["text"][:40] for c in cases]


ACUTE = [c for c in CASES if c["label"] == "acute"]
NOT_ACUTE = [c for c in CASES if c["label"] != "acute"]


def test_eval_set_size_and_labels() -> None:
    assert len(CASES) >= 60
    assert {c["label"] for c in CASES} == set(LABELS)
    assert len(ACUTE) >= 30


@pytest.mark.parametrize("case", ACUTE, ids=_ids(ACUTE))
def test_acute_recall_is_100_percent(case: dict) -> None:
    assert check(ScrubbedText(case["text"])).level == RiskLevel.ACUTE


@pytest.mark.parametrize("case", ACUTE, ids=_ids(ACUTE))
def test_acute_recall_survives_scrub(case: dict) -> None:
    # Real pipeline: scrub() runs first, so PII placeholders must not break crisis detection.
    assert check(scrub(case["text"])).level == RiskLevel.ACUTE


@pytest.mark.parametrize("case", NOT_ACUTE, ids=_ids(NOT_ACUTE))
def test_non_acute_cases_match_label(case: dict) -> None:
    # Exact level documents the calibration; the hard requirement is "never ACUTE" for these.
    level = check(ScrubbedText(case["text"])).level
    assert level != RiskLevel.ACUTE
    assert level == LABELS[case["label"]]


@pytest.mark.parametrize(
    "text",
    ["всё бессмысленно, устал от всего", "от таблеток только хуже, брошу", "не могу уснуть без пива"],
)
def test_demo_persona_phrases(text: str) -> None:
    level = check(ScrubbedText(text)).level
    if "бессмысленно" in text:
        assert level >= RiskLevel.POSSIBLE
    else:
        assert level == RiskLevel.NONE


@pytest.mark.parametrize(
    "text",
    ["не хочу умирать", "мыслей о суициде нет", "не думаю о самоубийстве", "суицидальных мыслей не было"],
)
def test_negation_downgrades_to_possible_not_none(text: str) -> None:
    result = check(ScrubbedText(text))
    assert result.level == RiskLevel.POSSIBLE
    assert result.negated_rules


def test_denied_mention_does_not_hide_affirmed_one() -> None:
    assert check(ScrubbedText("не думаю о суициде. думаю о суициде каждую ночь")).level == RiskLevel.ACUTE
    assert check(ScrubbedText("мысли о суициде не отпускают")).level == RiskLevel.ACUTE


def test_weapon_alone_possible_with_intent_acute() -> None:
    alone = check(ScrubbedText("в сейфе лежит наградное оружие"))
    assert alone.level == RiskLevel.POSSIBLE
    assert alone.categories == ["weapon_access"]
    combo = check(ScrubbedText("в сейфе лежит наградное оружие, я обуза"))
    assert combo.level == RiskLevel.ACUTE
    assert "combo.weapon_with_risk" in combo.matched_rules


@pytest.mark.parametrize(
    ("text", "llm", "level", "source"),
    [
        ("сегодня всё нормально", None, RiskLevel.NONE, "none"),
        ("сегодня всё нормально", RiskLevel.NONE, RiskLevel.NONE, "none"),
        ("сегодня всё нормально", RiskLevel.POSSIBLE, RiskLevel.POSSIBLE, "llm"),
        ("сегодня всё нормально", RiskLevel.ACUTE, RiskLevel.ACUTE, "llm"),
        ("я обуза", None, RiskLevel.POSSIBLE, "lexicon"),
        ("я обуза", RiskLevel.ACUTE, RiskLevel.ACUTE, "both"),
        ("хочу умереть", RiskLevel.NONE, RiskLevel.ACUTE, "lexicon"),  # LLM can never lower the risk
        ("хочу умереть", RiskLevel.POSSIBLE, RiskLevel.ACUTE, "both"),
    ],
)
def test_llm_signal_max_combination(text: str, llm: RiskLevel | None, level: RiskLevel, source: str) -> None:
    result = check(ScrubbedText(text), llm_signal=llm)
    assert result.level == level
    assert result.source == source


def test_result_contains_no_raw_text() -> None:
    text = "Алексей пишет: хочу покончить с собой сегодня вечером"
    result = check(ScrubbedText(text))
    dumped = result.model_dump_json()
    for word in ("покончить", "собой", "Алексей", "вечером", normalise(text)):
        assert word not in dumped
    fields = {"level", "categories", "matched_rules", "negated_rules", "source"}
    assert set(SafetyResult.model_fields) == fields
    with pytest.raises(Exception):  # noqa: B017 - frozen pydantic model raises ValidationError
        result.level = RiskLevel.NONE  # type: ignore[misc]


def test_empty_and_punctuation_only() -> None:
    assert check(ScrubbedText("")).level == RiskLevel.NONE
    assert check(ScrubbedText("?!...")).level == RiskLevel.NONE


def test_normalisation() -> None:
    assert normalise("Всё  БЕССМЫСЛЕННО!!!") == "все бессмысленно"
    assert normalise("жиииить") == "жить"
    assert normalise("сyицид") == "суицид"  # Latin "y" inside a Russian word
    assert normalise("pokonchit s soboy") == "покончит с собой"


def test_lexicon_has_all_categories_and_compiles() -> None:
    lex = load_lexicon()
    cats = {r.category for r in lex.rules}
    expected = {"suicide", "self_harm", "harm_others", "weapon_access", "domestic_violence", "hopelessness"}
    assert cats == expected
    assert "психиатр" in lex.status


def test_performance_under_5ms_per_message() -> None:
    texts = [c["text"] for c in CASES] + ["Сегодня был тяжёлый день, спал плохо, снились кошмары. " * 8]
    check(ScrubbedText("прогрев"))
    start = time.perf_counter()
    rounds = 20
    for _ in range(rounds):
        for t in texts:
            check(ScrubbedText(t))
    per_msg_ms = (time.perf_counter() - start) * 1000 / (rounds * len(texts))
    assert per_msg_ms < 5


def test_crisis_reply_patient_default() -> None:
    script = crisis_reply(["suicide"])
    assert isinstance(script, CrisisScript)
    assert script.variant == "patient.default"
    assert "112" in script.text
    assert "8-800-2000-122" in script.text
    assert "бесплатно, круглосуточно" in script.text
    assert "уведомлён" in script.text
    assert script.grounding in script.text
    assert "психиатр (#8)" in script.status


def test_crisis_reply_variants() -> None:
    dv = crisis_reply(["domestic_violence", "suicide"], audience="family")
    assert dv.variant == "family.domestic_violence"
    assert "потерпеть" in dv.text and "112" in dv.text and "8-800-2000-122" in dv.text
    assert crisis_reply(["domestic_violence"]).variant == "patient.domestic_violence"
    assert crisis_reply(["harm_others"]).variant == "patient.harm_others"
    assert crisis_reply(["harm_others"], audience="family").variant == "family.default"
    assert crisis_reply([], audience="family").variant == "family.default"