polyguard-rl/tests/test_env_step_flow.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_completed_episode_does_not_advance_on_extra_step() -> None:
    env = PolyGuardEnv()
    env.reset(seed=42, difficulty="medium", sub_environment="REGIMEN_RISK")
    env.state.max_steps = 1
    action = env.get_legal_actions()[0]

    _, _, done, _ = env.step(action)
    step_count = env.state.step_count
    _, _, repeated_done, info = env.step(action)

    assert done
    assert repeated_done
    assert env.state.step_count == step_count
    assert info["termination_reason"] == "already_done"


def test_minimized_dose_reduction_is_not_legal() -> None:
    env = PolyGuardEnv()
    env.reset(seed=42, difficulty="medium", sub_environment="REGIMEN_RISK")
    target = env.state.patient.medications[0].drug
    env.state.patient.medications[0].dose_bucket = DoseBucket.LOW
    action = PolyGuardAction(
        mode=DecisionMode.DOSE_OPT,
        action_type=ActionType.REDUCE_DOSE_BUCKET,
        target_drug=target,
        dose_bucket=DoseBucket.LOW,
        candidate_id="cand_03",
        confidence=0.7,
        rationale_brief="test",
    )

    report = verify_action_legality(env.state, action)

    assert not report.legal
    assert "dose_already_minimized" in report.violations


def test_review_request_does_not_end_episode_with_budget_remaining() -> None:
    env = PolyGuardEnv()
    env.reset(seed=42, difficulty="medium", sub_environment="REGIMEN_RISK")
    review_action = next(
        action for action in env.get_legal_actions() if action["action_type"] == ActionType.REQUEST_PHARMACIST_REVIEW.value
    )

    _, _, done, info = env.step(review_action)

    assert not done
    assert info["termination_reason"] == "ongoing"
    assert env.state.step_count < env.state.max_steps