Spaces:
Sleeping
Sleeping
File size: 5,825 Bytes
0cb452d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | """Quick integration test for the TriageFlow environment."""
import sys
sys.path.insert(0, '.')
from triage_flow.environment import TriageEnvironment
from triage_flow.graders import grade_task
from models import TriageAction, ActionType, PriorityLevel, InfoField
def test_basic_triage():
env = TriageEnvironment()
obs = env.reset(task_name="basic-triage")
pid = obs.current_patient["patient_id"] if obs.current_patient else "none"
print(f"Reset OK - patients: {obs.queue_length}, current: {pid}")
# Correct assignments for all 3 patients
for pid, p in [("P001", PriorityLevel.IMMEDIATE), ("P002", PriorityLevel.NON_URGENT), ("P003", PriorityLevel.URGENT)]:
action = TriageAction(action_type=ActionType.ASSIGN_PRIORITY, patient_id=pid, priority_level=p)
obs = env.step(action)
print(f" Assigned {pid}={p.value} -> reward={obs.reward}, done={obs.done}")
state = env.state
print(f"State: assignments={state.assignments}, cleared={state.queue_cleared}")
score = grade_task("basic-triage", state.model_dump())
print(f"Grader score: {score}")
assert score == 1.0, f"Expected 1.0, got {score}"
print("BASIC TRIAGE: PASSED\n")
def test_wrong_assignments():
env = TriageEnvironment()
env.reset(task_name="basic-triage")
for pid, p in [("P001", PriorityLevel.NON_URGENT), ("P002", PriorityLevel.IMMEDIATE), ("P003", PriorityLevel.NON_URGENT)]:
env.step(TriageAction(action_type=ActionType.ASSIGN_PRIORITY, patient_id=pid, priority_level=p))
score = grade_task("basic-triage", env.state.model_dump())
print(f"All wrong score: {score}")
assert score == 0.0, f"Expected 0.0, got {score}"
print("WRONG ASSIGNMENTS: PASSED\n")
def test_incomplete_records():
env = TriageEnvironment()
obs = env.reset(task_name="incomplete-records-triage")
pid = obs.current_patient["patient_id"] if obs.current_patient else "none"
print(f"Incomplete records - patients: {obs.queue_length}, current: {pid}")
print(f" Missing fields: {obs.missing_fields}")
# Request info for P101 (vitals missing)
action = TriageAction(action_type=ActionType.REQUEST_INFO, patient_id="P101", info_field=InfoField.VITALS)
obs = env.step(action)
print(f" Info request -> reward={obs.reward}")
# Assign correct priorities
for pid, p in [("P101", PriorityLevel.IMMEDIATE), ("P102", PriorityLevel.LESS_URGENT),
("P103", PriorityLevel.URGENT), ("P104", PriorityLevel.NON_URGENT)]:
obs = env.step(TriageAction(action_type=ActionType.ASSIGN_PRIORITY, patient_id=pid, priority_level=p))
score = grade_task("incomplete-records-triage", env.state.model_dump())
print(f"Incomplete records score: {score}")
assert 0.5 <= score <= 1.0, f"Expected good score, got {score}"
print("INCOMPLETE RECORDS: PASSED\n")
def test_mass_casualty():
env = TriageEnvironment()
obs = env.reset(task_name="mass-casualty-triage")
print(f"Mass casualty - patients: {obs.queue_length}")
# Assign all 8 patients with correct priorities
correct = [
("P201", PriorityLevel.IMMEDIATE), ("P202", PriorityLevel.IMMEDIATE),
("P203", PriorityLevel.LESS_URGENT), ("P204", PriorityLevel.LESS_URGENT),
("P205", PriorityLevel.URGENT), ("P206", PriorityLevel.NON_URGENT),
("P207", PriorityLevel.URGENT), ("P208", PriorityLevel.NON_URGENT),
]
for pid, p in correct:
obs = env.step(TriageAction(action_type=ActionType.ASSIGN_PRIORITY, patient_id=pid, priority_level=p))
score = grade_task("mass-casualty-triage", env.state.model_dump())
print(f"Mass casualty perfect score: {score}")
assert score >= 0.8, f"Expected high score, got {score}"
print("MASS CASUALTY: PASSED\n")
def test_grader_variance():
"""Different behaviors must produce different scores."""
scores = set()
# Perfect
env = TriageEnvironment()
env.reset(task_name="basic-triage")
for pid, p in [("P001", PriorityLevel.IMMEDIATE), ("P002", PriorityLevel.NON_URGENT), ("P003", PriorityLevel.URGENT)]:
env.step(TriageAction(action_type=ActionType.ASSIGN_PRIORITY, patient_id=pid, priority_level=p))
scores.add(grade_task("basic-triage", env.state.model_dump()))
# All wrong
env = TriageEnvironment()
env.reset(task_name="basic-triage")
for pid, p in [("P001", PriorityLevel.NON_URGENT), ("P002", PriorityLevel.IMMEDIATE), ("P003", PriorityLevel.NON_URGENT)]:
env.step(TriageAction(action_type=ActionType.ASSIGN_PRIORITY, patient_id=pid, priority_level=p))
scores.add(grade_task("basic-triage", env.state.model_dump()))
# Partial
env = TriageEnvironment()
env.reset(task_name="basic-triage")
for pid, p in [("P001", PriorityLevel.IMMEDIATE), ("P002", PriorityLevel.IMMEDIATE), ("P003", PriorityLevel.IMMEDIATE)]:
env.step(TriageAction(action_type=ActionType.ASSIGN_PRIORITY, patient_id=pid, priority_level=p))
scores.add(grade_task("basic-triage", env.state.model_dump()))
# No actions
env = TriageEnvironment()
env.reset(task_name="basic-triage")
scores.add(grade_task("basic-triage", env.state.model_dump()))
print(f"Grader variance - unique scores: {scores}")
assert len(scores) >= 3, f"Expected 3+ different scores, got {scores}"
print("GRADER VARIANCE: PASSED\n")
def test_server_app_import():
"""Test that the server app can be imported."""
from server.app import app
print(f"Server app imported: {app.title}")
print("SERVER IMPORT: PASSED\n")
if __name__ == "__main__":
test_basic_triage()
test_wrong_assignments()
test_incomplete_records()
test_mass_casualty()
test_grader_variance()
test_server_app_import()
print("=" * 60)
print("ALL INTEGRATION TESTS PASSED!")
print("=" * 60)
|