File size: 5,679 Bytes
4a0825f | 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 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 | import pytest
from langchain_core.messages import HumanMessage, AIMessage, ToolMessage
from src.core.state import AgentState
from src.core.graph import (
detect_emergency,
role_classifier_node,
emergency_response_node,
recovery_loop_node,
route_after_role,
route_patient_llm,
route_after_validator,
route_after_safety,
route_after_intent,
medical_pipeline
)
from src.core.graph_cdm import (
data_fetch_node,
trend_analyzer_node,
cdm_pipeline
)
def test_detect_emergency():
# Healthy case
state = {"messages": [HumanMessage(content="Hello doctor, I have a normal question about my diet.")]}
assert not detect_emergency(state)
# Emergency cases
emergencies = [
"My grandfather just passed out and is unconscious!",
"He is having a seizure right now.",
"I think I am having a stroke or heart attack, I have chest pain."
]
for msg in emergencies:
state = {"messages": [HumanMessage(content=msg)]}
assert detect_emergency(state) is True
@pytest.mark.asyncio
async def test_role_classifier_node():
# User message indicating caregiver
state = {
"messages": [HumanMessage(content="I am caring for my mother who has diabetes.")],
"user_role": "unknown",
"intent_type": "unknown"
}
res = await role_classifier_node(state)
assert res["user_role"] == "caregiver"
assert "logs" in res
@pytest.mark.asyncio
async def test_emergency_response_node():
state = {"messages": [HumanMessage(content="Heart attack!")], "logs": []}
res = await emergency_response_node(state)
assert res["is_valid"] is True
assert res["is_safe"] is True
assert len(res["messages"]) == 1
assert "medical emergency" in res["messages"][0].content
@pytest.mark.asyncio
async def test_recovery_loop_node():
state = {"attempts": 1, "logs": []}
res = await recovery_loop_node(state)
assert res["attempts"] == 2
assert "[RECOVERY]" in res["messages"][0].content
def test_routing_logic():
# route_after_role
state = {"user_role": "patient", "messages": [HumanMessage(content="Normal checkup")]}
assert route_after_role(state) == "patient_llm"
state_emerg = {"user_role": "patient", "messages": [HumanMessage(content="passed out")]}
assert route_after_role(state_emerg) == "emergency_response"
state_clin = {"user_role": "clinician"}
assert route_after_role(state_clin) == "intent_classifier"
# route_patient_llm
state_tools = {"messages": [AIMessage(content="Need tools", tool_calls=[{"name": "web_search", "args": {}, "id": "tc1"}])]}
assert route_patient_llm(state_tools) == "tools_node"
state_validator = {"messages": [AIMessage(content="Normal response")]}
assert route_patient_llm(state_validator) == "validator"
# route_after_validator
assert route_after_validator({"is_valid": True}) == "safety_check"
assert route_after_validator({"is_valid": False}) == "recovery_loop"
# route_after_safety
assert route_after_safety({"is_safe": True}) == "persistence_node"
assert route_after_safety({"is_safe": False}) == "recovery_loop"
# route_after_intent
assert route_after_intent({"intent_type": "diagnosis"}) == "diagnosis_assist"
assert route_after_intent({"intent_type": "treatment"}) == "treatment_assist"
assert route_after_intent({"intent_type": "monitoring"}) == "monitoring_assist"
assert route_after_intent({"intent_type": "other"}) == "general_assist"
@pytest.mark.asyncio
async def test_medical_pipeline_e2e(clean_db):
patient_id = "patient-graph-test"
clean_db._ensure_patient_exists(patient_id, clean_db)
initial_state = {
"messages": [HumanMessage(content="Hello assistant, I need general help.")],
"user_role": "unknown",
"intent_type": "unknown",
"session_id": "session-graph-1",
"patient_id": patient_id,
"is_valid": False,
"is_safe": False,
"attempts": 0,
"clinician_outputs": [],
"patient_response": "",
"research_output": "",
"sources": [],
"logs": [],
"metrics": [],
}
final_state = await medical_pipeline.ainvoke(
initial_state,
config={"recursion_limit": 50}
)
assert final_state["user_role"] == "patient"
assert final_state["is_valid"] is True
assert final_state["is_safe"] is True
assert len(final_state["messages"]) > 1
@pytest.mark.asyncio
async def test_cdm_pipeline_e2e(clean_db):
patient_id = "patient-cdm-graph"
# Seed data
clean_db._ensure_patient_exists(patient_id, clean_db)
clean_db._insert("observations", {
"id": "obs-1",
"patient_id": patient_id,
"resource": {
"resourceType": "Observation",
"code": {"coding": [{"display": "Glucose"}]},
"valueQuantity": {"value": 150}
},
"effectiveDateTime": "2026-06-05T12:00:00Z"
})
initial_state = {
"messages": [HumanMessage(content="Let's analyze trends.")],
"user_role": "patient",
"patient_id": patient_id,
"session_id": "session-cdm-1",
"is_valid": False,
"is_safe": False,
"attempts": 0,
"clinician_outputs": [],
"patient_response": "",
"research_output": "",
"sources": [],
"logs": [],
"metrics": [],
}
final_state = await cdm_pipeline.ainvoke(
initial_state,
config={"recursion_limit": 50}
)
assert "trend_analysis" in final_state
assert final_state["is_valid"] is True
assert final_state["is_safe"] is True
|