polyguard-rl/tests/test_timeout_logic.py

from app.common.enums import ActionType, DecisionMode, Difficulty, DoseBucket, SubEnvironment
from app.common.types import LabSummary, Medication, PatientProfile, PolyGuardAction, PolyGuardState
from app.env.termination import check_termination_with_timeout


def _state() -> PolyGuardState:
    patient = PatientProfile(
        patient_id="p1",
        age=65,
        sex="F",
        medications=[Medication(drug="warfarin_like")],
        labs=LabSummary(),
        vitals={},
    )
    return PolyGuardState(
        episode_id="ep1",
        seed=1,
        scenario_id="s1",
        difficulty=Difficulty.EASY,
        sub_environment=SubEnvironment.DDI,
        step_count=0,
        max_steps=3,
        patient=patient,
    )


def _action() -> PolyGuardAction:
    return PolyGuardAction(
        mode=DecisionMode.REGIMEN_OPT,
        action_type=ActionType.KEEP_REGIMEN,
        target_drug=None,
        replacement_drug=None,
        dose_bucket=DoseBucket.NA,
        taper_days=None,
        monitoring_plan=None,
        candidate_id="cand_01",
        confidence=0.7,
        rationale_brief="test",
    )


def test_wall_clock_timeout_trigger() -> None:
    done, reason = check_termination_with_timeout(
        state=_state(),
        action=_action(),
        elapsed_seconds=5.0,
        wall_clock_limit_seconds=1.0,
    )
    assert done
    assert reason == "wall_clock_timeout"