aionboarding/ai-patient-onboarding/tests/test_websocket.py

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