polyguard-rl/tests/test_constraints.py

from app.common.enums import ActionType, DecisionMode, DoseBucket
from app.common.types import PolyGuardAction
from app.env.env_core import PolyGuardEnv
from app.env.verifier import verify_action_legality


def test_abrupt_stop_requires_taper() -> None:
    env = PolyGuardEnv()
    env.reset(seed=42, difficulty="easy")
    action = PolyGuardAction(
        mode=DecisionMode.REGIMEN_OPT,
        action_type=ActionType.STOP_DRUG,
        target_drug="benzodiazepine_like",
        replacement_drug=None,
        dose_bucket=DoseBucket.NA,
        taper_days=None,
        monitoring_plan=None,
        candidate_id="cand_99",
        confidence=0.8,
        rationale_brief="stop",
    )
    report = verify_action_legality(env.state, action)
    assert report.legal is False
    assert "abrupt_stop_requires_taper" in report.violations


def test_monitoring_action_requires_plan() -> None:
    env = PolyGuardEnv()
    env.reset(seed=42, difficulty="medium")
    action = PolyGuardAction(
        mode=DecisionMode.DOSE_OPT,
        action_type=ActionType.ORDER_MONITORING_AND_WAIT,
        target_drug=None,
        replacement_drug=None,
        dose_bucket=DoseBucket.NA,
        taper_days=None,
        monitoring_plan=None,
        candidate_id="cand_08",
        confidence=0.7,
        rationale_brief="wait",
    )
    report = verify_action_legality(env.state, action)
    assert report.legal is False
    assert "invalid_monitoring_gap" in report.violations