test_redesign.py

"""Quick validation of the redesigned environment lifecycle."""
import sys, os
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from triage_flow.environment import TriageEnvironment
from triage_flow.graders import grade_task
from models import TriageAction

env = TriageEnvironment()

# ---- Task 1: Perfect run ----
obs = env.reset(task_name="basic-triage")
print(f"Step 0: patient={obs.incoming_patient['patient_id']}, queue={obs.current_queue}")

action = TriageAction(classification="immediate", reordered_queue=["P001"])
obs = env.step(action)
print(f"Step 1: reward={obs.reward}, queue={obs.current_queue}, done={obs.done}")

action = TriageAction(classification="non_urgent", reordered_queue=["P001", "P002"])
obs = env.step(action)
print(f"Step 2: reward={obs.reward}, queue={obs.current_queue}, done={obs.done}")

action = TriageAction(classification="urgent", reordered_queue=["P001", "P003", "P002"])
obs = env.step(action)
print(f"Step 3: reward={obs.reward}, queue={obs.current_queue}, done={obs.done}")

score = grade_task("basic-triage", env.state.model_dump())
print(f"PERFECT RUN SCORE: {score}")

# ---- Task 1: Imperfect run (all less_urgent) ----
env2 = TriageEnvironment()
obs = env2.reset(task_name="basic-triage")

for i in range(3):
    pid = obs.incoming_patient["patient_id"]
    q = list(obs.current_queue) + [pid]
    action = TriageAction(classification="less_urgent", reordered_queue=q)
    obs = env2.step(action)

score2 = grade_task("basic-triage", env2.state.model_dump())
print(f"IMPERFECT RUN SCORE: {score2}")
print(f"Scores differ: {score != score2}")

# ---- Task 2: Test escalation ----
env3 = TriageEnvironment()
obs = env3.reset(task_name="incomplete-records-triage")
print(f"\nTask 2 first patient: {obs.incoming_patient['patient_id']}, info_complete={obs.incoming_patient.get('info_complete')}")

# Correctly escalate P101
action = TriageAction(classification="escalate", reordered_queue=[])
obs = env3.step(action)
print(f"After escalation: reward={obs.reward}, queue={obs.current_queue}")