| 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(): |
| |
| state = {"messages": [HumanMessage(content="Hello doctor, I have a normal question about my diet.")]} |
| assert not detect_emergency(state) |
|
|
| |
| 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(): |
| |
| 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(): |
| |
| 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" |
|
|
| |
| 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" |
|
|
| |
| assert route_after_validator({"is_valid": True}) == "safety_check" |
| assert route_after_validator({"is_valid": False}) == "recovery_loop" |
|
|
| |
| assert route_after_safety({"is_safe": True}) == "persistence_node" |
| assert route_after_safety({"is_safe": False}) == "recovery_loop" |
|
|
| |
| 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" |
| |
| 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 |
|
|
