""" 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