295 lines
9.5 KiB
Python
295 lines
9.5 KiB
Python
import uuid
|
|
import asyncio
|
|
import concurrent.futures
|
|
import pytest
|
|
from unittest.mock import AsyncMock, patch, MagicMock
|
|
|
|
from app.ai.gemini_client import GeminiClient
|
|
from app.ai.sarvam_client import (
|
|
SarvamClient,
|
|
STT_TIMEOUT,
|
|
STT_NETWORK_ERROR,
|
|
STT_SUCCESS,
|
|
TTS_TIMEOUT,
|
|
TTS_NETWORK_ERROR,
|
|
TTS_SUCCESS,
|
|
)
|
|
from app.database.database import SessionLocal, engine, Base
|
|
from app.database.models import OnboardingSession, OnboardingState, OnboardingMessage
|
|
|
|
|
|
# =====================================================================
|
|
# GEMINI CLIENT TESTS
|
|
# =====================================================================
|
|
|
|
@pytest.mark.anyio
|
|
async def test_gemini_client():
|
|
client = GeminiClient()
|
|
assert client.model is not None
|
|
# Unconfigured key fallback
|
|
with patch.object(client, "api_key", ""):
|
|
doc_res = await client.extract_from_document(b"fake", "application/pdf", {})
|
|
assert doc_res["document_type"] == "medical_report"
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_gemini_timeout():
|
|
client = GeminiClient()
|
|
client.api_key = "valid_fake_key"
|
|
with patch("asyncio.wait_for", side_effect=asyncio.TimeoutError):
|
|
with pytest.raises(Exception):
|
|
await client.generate_response("System prompt", "User msg")
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_gemini_provider_error():
|
|
client = GeminiClient()
|
|
client.api_key = "valid_fake_key"
|
|
with patch("asyncio.wait_for", side_effect=RuntimeError("API quota exceeded")):
|
|
with pytest.raises(Exception):
|
|
await client.generate_response("System prompt", "User msg")
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_gemini_malformed_response():
|
|
client = GeminiClient()
|
|
client.api_key = "valid_fake_key"
|
|
mock_resp = MagicMock()
|
|
mock_resp.text = "NOT JSON VALUE"
|
|
with patch("asyncio.wait_for", return_value=mock_resp):
|
|
res = await client.extract_from_document(b"data", "image/png", {})
|
|
assert res["document_type"] == "unknown"
|
|
|
|
|
|
# =====================================================================
|
|
# VOICE STT / TTS RELIABILITY TESTS
|
|
# =====================================================================
|
|
|
|
@pytest.mark.anyio
|
|
async def test_stt_timeout():
|
|
sarvam = SarvamClient()
|
|
sarvam.api_key = "fake_key"
|
|
with patch.object(sarvam, "_get_client") as mock_get_client:
|
|
mock_client = AsyncMock()
|
|
mock_client.post.side_effect = asyncio.TimeoutError
|
|
mock_get_client.return_value = mock_client
|
|
transcript, status = await sarvam.speech_to_text_with_status(b"audio_bytes", max_retries=1)
|
|
assert transcript == ""
|
|
assert status == STT_TIMEOUT
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_stt_network_failure():
|
|
sarvam = SarvamClient()
|
|
sarvam.api_key = "fake_key"
|
|
with patch.object(sarvam, "_get_client") as mock_get_client:
|
|
mock_client = AsyncMock()
|
|
mock_client.post.side_effect = Exception("Network unreachable")
|
|
mock_get_client.return_value = mock_client
|
|
transcript, status = await sarvam.speech_to_text_with_status(b"audio_bytes", max_retries=1)
|
|
assert transcript == ""
|
|
assert status == STT_NETWORK_ERROR
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_stt_retry_limit():
|
|
sarvam = SarvamClient()
|
|
sarvam.api_key = "fake_key"
|
|
attempt_count = 0
|
|
async def mock_post(*args, **kwargs):
|
|
nonlocal attempt_count
|
|
attempt_count += 1
|
|
raise asyncio.TimeoutError
|
|
|
|
with patch.object(sarvam, "_get_client") as mock_get_client:
|
|
mock_client = AsyncMock()
|
|
mock_client.post = mock_post
|
|
mock_get_client.return_value = mock_client
|
|
await sarvam.speech_to_text_with_status(b"audio_bytes", max_retries=2)
|
|
assert attempt_count == 3 # Initial + 2 retries
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_tts_timeout():
|
|
sarvam = SarvamClient()
|
|
sarvam.api_key = "fake_key"
|
|
with patch.object(sarvam, "_get_client") as mock_get_client:
|
|
mock_client = AsyncMock()
|
|
mock_client.post.side_effect = asyncio.TimeoutError
|
|
mock_get_client.return_value = mock_client
|
|
audio, status = await sarvam.text_to_speech_with_status("Hello", max_retries=1)
|
|
assert audio is None
|
|
assert status == TTS_TIMEOUT
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_tts_network_failure():
|
|
sarvam = SarvamClient()
|
|
sarvam.api_key = "fake_key"
|
|
with patch.object(sarvam, "_get_client") as mock_get_client:
|
|
mock_client = AsyncMock()
|
|
mock_client.post.side_effect = Exception("Network unreachable")
|
|
mock_get_client.return_value = mock_client
|
|
audio, status = await sarvam.text_to_speech_with_status("Hello", max_retries=1)
|
|
assert audio is None
|
|
assert status == TTS_NETWORK_ERROR
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_tts_retry_limit():
|
|
sarvam = SarvamClient()
|
|
sarvam.api_key = "fake_key"
|
|
attempt_count = 0
|
|
async def mock_post(*args, **kwargs):
|
|
nonlocal attempt_count
|
|
attempt_count += 1
|
|
raise asyncio.TimeoutError
|
|
|
|
with patch.object(sarvam, "_get_client") as mock_get_client:
|
|
mock_client = AsyncMock()
|
|
mock_client.post = mock_post
|
|
mock_get_client.return_value = mock_client
|
|
await sarvam.text_to_speech_with_status("Hello", max_retries=2)
|
|
assert attempt_count == 3
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_voice_failure_does_not_crash_websocket():
|
|
sarvam = SarvamClient()
|
|
with patch.object(sarvam, "speech_to_text_with_status", return_value=("", STT_TIMEOUT)):
|
|
transcript, status = await sarvam.speech_to_text_with_status(b"audio", max_retries=0)
|
|
assert status == STT_TIMEOUT
|
|
assert transcript == ""
|
|
|
|
|
|
# =====================================================================
|
|
# SQLITE CONCURRENCY TESTS
|
|
# =====================================================================
|
|
|
|
def _worker_write_session(session_idx: int):
|
|
db = SessionLocal()
|
|
try:
|
|
session_id = f"CONC-SESS-{uuid.uuid4()}"
|
|
sess = OnboardingSession(
|
|
session_id=session_id,
|
|
patient_id=f"PAT-{session_idx}",
|
|
language="en-IN",
|
|
category_type="general",
|
|
category_code="general",
|
|
category_name="General",
|
|
stage="chief_complaint",
|
|
status="in_progress",
|
|
)
|
|
db.add(sess)
|
|
db.commit()
|
|
|
|
msg = OnboardingMessage(
|
|
session_id=session_id,
|
|
role="patient",
|
|
message="Severe fever for 2 days",
|
|
stage="chief_complaint",
|
|
)
|
|
db.add(msg)
|
|
db.commit()
|
|
return True
|
|
except Exception as e:
|
|
db.rollback()
|
|
raise e
|
|
finally:
|
|
db.close()
|
|
|
|
|
|
def test_concurrent_session_writes():
|
|
"""Simulate 15 concurrent session writes to SQLite under WAL mode."""
|
|
Base.metadata.create_all(bind=engine)
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
|
|
futures = [executor.submit(_worker_write_session, i) for i in range(15)]
|
|
results = [f.result() for f in concurrent.futures.as_completed(futures)]
|
|
assert all(results)
|
|
|
|
|
|
def test_concurrent_websocket_sessions():
|
|
"""Simulate multiple concurrent session updates without database locks."""
|
|
Base.metadata.create_all(bind=engine)
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
|
|
futures = [executor.submit(_worker_write_session, i + 100) for i in range(5)]
|
|
results = [f.result() for f in concurrent.futures.as_completed(futures)]
|
|
assert len(results) == 5
|
|
|
|
|
|
def test_database_lock_recovery():
|
|
db1 = SessionLocal()
|
|
db2 = SessionLocal()
|
|
try:
|
|
s_id = f"LOCK-TEST-{uuid.uuid4()}"
|
|
sess1 = OnboardingSession(
|
|
session_id=s_id,
|
|
patient_id="PAT-L1",
|
|
language="en",
|
|
category_type="general",
|
|
category_code="general",
|
|
category_name="General",
|
|
stage="chief_complaint",
|
|
status="in_progress",
|
|
)
|
|
db1.add(sess1)
|
|
db1.commit()
|
|
|
|
res1 = db1.query(OnboardingSession).filter_by(session_id=s_id).first()
|
|
res2 = db2.query(OnboardingSession).filter_by(session_id=s_id).first()
|
|
assert res1 is not None
|
|
assert res2 is not None
|
|
finally:
|
|
db1.close()
|
|
db2.close()
|
|
|
|
|
|
def test_transaction_rollback():
|
|
db = SessionLocal()
|
|
try:
|
|
s_id = f"ROLLBACK-TEST-{uuid.uuid4()}"
|
|
sess = OnboardingSession(
|
|
session_id=s_id,
|
|
patient_id="PAT-R1",
|
|
language="en",
|
|
category_type="general",
|
|
category_code="general",
|
|
category_name="General",
|
|
stage="chief_complaint",
|
|
status="in_progress",
|
|
)
|
|
db.add(sess)
|
|
db.flush()
|
|
db.rollback()
|
|
check = db.query(OnboardingSession).filter_by(session_id=s_id).first()
|
|
assert check is None
|
|
finally:
|
|
db.close()
|
|
|
|
|
|
def test_disconnect_during_db_operation():
|
|
db = SessionLocal()
|
|
try:
|
|
s_id = f"DISCONNECT-TEST-{uuid.uuid4()}"
|
|
sess = OnboardingSession(
|
|
session_id=s_id,
|
|
patient_id="PAT-D1",
|
|
language="en",
|
|
category_type="general",
|
|
category_code="general",
|
|
category_name="General",
|
|
stage="chief_complaint",
|
|
status="in_progress",
|
|
)
|
|
db.add(sess)
|
|
db.commit()
|
|
db.close()
|
|
|
|
db_new = SessionLocal()
|
|
check = db_new.query(OnboardingSession).filter_by(session_id=s_id).first()
|
|
assert check is not None
|
|
db_new.close()
|
|
except Exception:
|
|
db.close()
|
|
raise
|