feat: complete production hardening implementation, specialty canonicalization, google-genai migration, voice reliability, sqlite concurrency, and test suite consolidation

ai-patient-onboarding/tests/test_websocket.py

@@ -0,0 +1,250 @@
+"""
+Comprehensive WebSocket Test Suite (Section 37, 38, 40).
+Tests connection lifecycle, handshake, message protocol, typing/processing events,
+deduplication, safety alerts, error protocol, multilingual handling, and REST backward compatibility.
+"""
+
+import sys
+import os
+import json
+import pytest
+from fastapi.testclient import TestClient
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+from app.main import app
+
+client = TestClient(app)
+
+
+def create_test_session(patient_id: str = "PAT-WS-100", language: str = "en") -> str:
+    """Helper to initialize an active onboarding session via REST API."""
+    response = client.post(
+        "/v1/onboarding/start",
+        json={
+            "patient": {"id": patient_id, "name": "WebSocket Test Patient", "age": 32, "gender": "female"},
+            "category": {"type": "general", "code": "general_medicine", "name": "General Medicine"},
+            "language": language
+        },
+        headers={"Authorization": f"Bearer test_token_{patient_id}"}
+    )
+    assert response.status_code == 200
+    return response.json()["session_id"]
+
+
+def test_websocket_successful_connection_ack():
+    """Test WS connection lifecycle and connection.ack event (Section 3 & 4)."""
+    patient_id = "PAT-WS-ACK"
+    session_id = create_test_session(patient_id=patient_id)
+
+    url = f"/api/v1/ws/onboarding/{session_id}?token=test_token_{patient_id}"
+    with client.websocket_connect(url) as websocket:
+        ack_data = websocket.receive_json()
+        assert ack_data["type"] == "connection.ack"
+        assert ack_data["session_id"] == session_id
+        assert ack_data["status"] == "connected"
+
+
+def test_websocket_invalid_session_connection():
+    """Test connection to non-existent session ID is rejected (Section 24)."""
+    url = "/api/v1/ws/onboarding/ONB-INVALID-9999?token=test_token_PAT-WS-1"
+    with client.websocket_connect(url) as websocket:
+        err = websocket.receive_json()
+        assert err["type"] == "error"
+        assert err["code"] == "SESSION_NOT_FOUND"
+
+
+def test_websocket_unauthorized_session_access():
+    """Test connecting to Session A with User B token is rejected (Section 25)."""
+    session_id = create_test_session(patient_id="PAT-USER-A")
+
+    # Connect with token for PAT-USER-B
+    url = f"/api/v1/ws/onboarding/{session_id}?token=test_token_PAT-USER-B"
+    with client.websocket_connect(url) as websocket:
+        err = websocket.receive_json()
+        assert err["type"] == "error"
+        assert err["code"] == "UNAUTHORIZED"
+
+
+def test_websocket_end_to_end_message_flow():
+    """Required End-to-End Test (Section 38): verify complete event sequence."""
+    patient_id = "PAT-WS-E2E"
+    session_id = create_test_session(patient_id=patient_id)
+
+    url = f"/api/v1/ws/onboarding/{session_id}?token=test_token_{patient_id}"
+    with client.websocket_connect(url) as websocket:
+        # Receive ACK
+        ack = websocket.receive_json()
+        assert ack["type"] == "connection.ack"
+
+        # Send patient.message
+        websocket.send_json({
+            "type": "patient.message",
+            "message_id": "msg_e2e_1",
+            "text": "I have severe headache for 3 days"
+        })
+
+        events = []
+        # Expect: processing.started, typing.started, state.updated, assistant.message, typing.stopped, processing.completed
+        for _ in range(6):
+            event = websocket.receive_json()
+            events.append(event)
+
+        event_types = [e["type"] for e in events]
+        assert "processing.started" in event_types
+        assert "typing.started" in event_types
+        assert "state.updated" in event_types
+        assert "assistant.message" in event_types
+        assert "typing.stopped" in event_types
+        assert "processing.completed" in event_types
+
+        # Verify assistant response properties (Section 5 & 9)
+        asst_msg = next(e for e in events if e["type"] == "assistant.message")
+        assert asst_msg["message_id"] == "resp_msg_e2e_1"
+        assert len(asst_msg["reply"].split()) <= 20
+        assert asst_msg["reply"].count("?") <= 1
+
+
+def test_websocket_duplicate_message_prevention():
+    """Test duplicate message_id protection returns cached response (Section 17)."""
+    patient_id = "PAT-WS-DEDUP"
+    session_id = create_test_session(patient_id=patient_id)
+
+    url = f"/api/v1/ws/onboarding/{session_id}?token=test_token_{patient_id}"
+    with client.websocket_connect(url) as websocket:
+        websocket.receive_json()  # ACK
+
+        msg_payload = {
+            "type": "patient.message",
+            "message_id": "msg_dedup_101",
+            "text": "Fever for two days"
+        }
+
+        # First send
+        websocket.send_json(msg_payload)
+        events_1 = [websocket.receive_json() for _ in range(6)]
+        asst_1 = next(e for e in events_1 if e["type"] == "assistant.message")
+
+        # Duplicate send with SAME message_id
+        websocket.send_json(msg_payload)
+        cached_resp = websocket.receive_json()
+        assert cached_resp["type"] == "assistant.message"
+        assert cached_resp["reply"] == asst_1["reply"]
+
+
+def test_websocket_emergency_safety_alert():
+    """Test emergency message triggers safety.alert event (Section 12)."""
+    patient_id = "PAT-WS-EMERGENCY"
+    session_id = create_test_session(patient_id=patient_id)
+
+    url = f"/api/v1/ws/onboarding/{session_id}?token=test_token_{patient_id}"
+    with client.websocket_connect(url) as websocket:
+        websocket.receive_json()  # ACK
+
+        websocket.send_json({
+            "type": "patient.message",
+            "message_id": "msg_emerg_1",
+            "text": "I have severe crushing chest pain radiating to my arm and shortness of breath"
+        })
+
+        received = []
+        for _ in range(8):  # Includes safety alert and onboarding completed
+            received.append(websocket.receive_json())
+
+        event_types = [e["type"] for e in received]
+        assert "safety.alert" in event_types
+
+        safety_event = next(e for e in received if e["type"] == "safety.alert")
+        assert safety_event["level"] in ["emergency", "crisis"]
+
+
+def test_websocket_invalid_json_handling():
+    """Test malformed non-JSON text returns error event without closing socket (Section 15)."""
+    patient_id = "PAT-WS-BADJSON"
+    session_id = create_test_session(patient_id=patient_id)
+
+    url = f"/api/v1/ws/onboarding/{session_id}?token=test_token_{patient_id}"
+    with client.websocket_connect(url) as websocket:
+        websocket.receive_json()  # ACK
+
+        # Send raw non-JSON text
+        websocket.send_text("THIS_IS_NOT_JSON")
+        err_resp = websocket.receive_json()
+        assert err_resp["type"] == "error"
+        assert err_resp["code"] == "INVALID_MESSAGE"
+
+        # Verify socket is still active by sending valid message after
+        websocket.send_json({
+            "type": "patient.message",
+            "message_id": "msg_valid_after_error",
+            "text": "I am fine now"
+        })
+        events = [websocket.receive_json() for _ in range(6)]
+        assert any(e["type"] == "assistant.message" for e in events)
+
+
+def test_websocket_oversized_message():
+    """Test message > 5000 chars returns MESSAGE_TOO_LARGE error (Section 16)."""
+    patient_id = "PAT-WS-LARGE"
+    session_id = create_test_session(patient_id=patient_id)
+
+    url = f"/api/v1/ws/onboarding/{session_id}?token=test_token_{patient_id}"
+    with client.websocket_connect(url) as websocket:
+        websocket.receive_json()  # ACK
+
+        oversized_text = "A" * 5050
+        websocket.send_json({
+            "type": "patient.message",
+            "message_id": "msg_large_1",
+            "text": oversized_text
+        })
+
+        err_resp = websocket.receive_json()
+        assert err_resp["type"] == "error"
+        assert err_resp["code"] == "MESSAGE_TOO_LARGE"
+
+
+def test_websocket_multilingual_unicode():
+    """Test Tamil/Hindi multilingual messages process seamlessly over WebSocket (Section 16)."""
+    patient_id = "PAT-WS-TAMIL"
+    session_id = create_test_session(patient_id=patient_id, language="ta")
+
+    url = f"/api/v1/ws/onboarding/{session_id}?token=test_token_{patient_id}"
+    with client.websocket_connect(url) as websocket:
+        websocket.receive_json()  # ACK
+
+        websocket.send_json({
+            "type": "patient.message",
+            "message_id": "msg_ta_1",
+            "text": "எனக்கு 3 நாளா தலைவலி இருக்கு"
+        })
+
+        events = [websocket.receive_json() for _ in range(6)]
+        asst = next(e for e in events if e["type"] == "assistant.message")
+        assert asst["reply"] is not None
+        assert len(asst["reply"]) > 0
+
+
+def test_rest_api_backward_compatibility():
+    """Verify Section 36 requirement: REST APIs continue working untouched."""
+    session_id = create_test_session(patient_id="PAT-REST-BACKWARD")
+
+    # Send REST message
+    res = client.post(
+        f"/v1/onboarding/{session_id}/message",
+        json={"message": "Headache for two days"},
+        headers={"Authorization": "Bearer test_token_PAT-REST-BACKWARD"}
+    )
+    assert res.status_code == 200
+    assert res.json()["success"] is True
+    assert "message" in res.json()
+
+
+def test_websocket_tester_ui_endpoint():
+    """Verify that the sample UI page for WebSocket testing returns 200 OK HTML."""
+    res = client.get("/ws-tester")
+    assert res.status_code == 200
+    assert "text/html" in res.headers["content-type"]
+    assert "AI Patient Onboarding" in res.text
+    assert "WebSocket" in res.text
+