"""Integration tests for POST /query endpoint.""" import asyncio from unittest.mock import AsyncMock, MagicMock, patch from uuid import uuid4 import pytest from fastapi.testclient import TestClient from openai.types.chat import ChatCompletion, ChatCompletionMessage from openai.types.chat.chat_completion import Choice from sqlalchemy import select from src.api.main import app from src.db.models import ApiKey, Document from src.db.session import get_sync_session from src.services.auth_service import hash_api_key @pytest.fixture def test_api_key(): """Create a test API key in the database.""" test_key = "test_api_key_for_query_12345" key_hash = hash_api_key(test_key) with get_sync_session() as session: # Clean up any existing test API key existing = session.execute( select(ApiKey).where(ApiKey.key_hash == key_hash) ).scalar_one_or_none() if existing: session.delete(existing) session.commit() # Create new test API key api_key = ApiKey( key_hash=key_hash, name="Test Query API Key", rate_limit_per_minute=60, is_active=True, ) session.add(api_key) session.commit() session.refresh(api_key) yield test_key, api_key.id # Cleanup: delete test API key if api_key_obj := session.execute( select(ApiKey).where(ApiKey.key_hash == key_hash) ).scalar_one_or_none(): session.delete(api_key_obj) session.commit() @pytest.fixture def test_document(): """Create a test completed document in the database.""" doc_id = uuid4() with get_sync_session() as session: # Create test document document = Document( id=doc_id, filename="test_document.xlsx", status="completed", file_size=1024, mime_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", ) session.add(document) session.commit() yield doc_id # Cleanup: delete test document if doc_obj := session.execute( select(Document).where(Document.id == doc_id) ).scalar_one_or_none(): session.delete(doc_obj) session.commit() @pytest.fixture def test_client(): """Create a FastAPI test client.""" return TestClient(app) @pytest.fixture def mock_milvus_search(): """Mock Milvus search results.""" with patch("src.retrieval.semantic_retriever.connections") as mock_conn: with patch("src.retrieval.semantic_retriever.Collection") as mock_coll_class: # Create mock collection mock_collection = MagicMock() mock_collection.load = MagicMock() # Create mock search result mock_hit = MagicMock() mock_hit.distance = 0.85 mock_hit.entity.get = lambda key, default="": { "filename": "test_document.xlsx", "sheet_name": "Sheet1", "col_range": "A1:D1", "text_content": "Item Code: 2024-4_0019, Status: Active", }.get(key, default) mock_collection.search.return_value = [[mock_hit]] mock_coll_class.return_value = mock_collection yield mock_collection @pytest.fixture def mock_azure_openai(): """Mock Azure OpenAI response.""" with patch("src.retrieval.answer_generator.AsyncAzureOpenAI") as mock_class: mock_client = MagicMock() # Create mock response mock_message = ChatCompletionMessage( role="assistant", content="Item code 2024-4_0019 has status Active as shown in Sheet1.", ) mock_choice = Choice(index=0, message=mock_message, finish_reason="stop") mock_completion = ChatCompletion( id="test-completion", model="gpt-4", object="chat.completion", created=1234567890, choices=[mock_choice], ) mock_client.chat.completions.create = AsyncMock(return_value=mock_completion) mock_class.return_value = mock_client yield mock_client @pytest.fixture def mock_embeddings(): """Mock embedding service.""" with patch("src.retrieval.semantic_retriever.EmbeddingService") as mock_class: mock_service = MagicMock() mock_service.embed_batch = AsyncMock(return_value=[[0.1] * 3072]) mock_class.return_value = mock_service yield mock_service class TestQueryEndpoint: """Tests for POST /query endpoint.""" def test_query_valid_request_returns_200( self, test_client, test_api_key, mock_milvus_search, mock_azure_openai, mock_embeddings, ): """Test successful query returns 200 with answer and sources.""" test_key, _ = test_api_key # Send query request response = test_client.post( "/query", json={ "query": "What is item code 2024-4_0019?", "include_context": True, }, headers={"X-API-Key": test_key}, ) # Verify response assert response.status_code == 200 data = response.json() assert "query" in data assert data["query"] == "What is item code 2024-4_0019?" assert "answer" in data assert "2024-4_0019" in data["answer"] assert "confidence" in data assert 0.0 <= data["confidence"] <= 1.0 assert "sources" in data assert len(data["sources"]) == 1 assert data["sources"][0]["file"] == "test_document.xlsx" assert "processing_time_ms" in data assert data["processing_time_ms"] > 0 def test_query_with_document_ids_filter( self, test_client, test_api_key, test_document, mock_milvus_search, mock_azure_openai, mock_embeddings, ): """Test query with document_ids filter.""" test_key, _ = test_api_key # Send query with document_ids filter response = test_client.post( "/query", json={ "query": "test query", "document_ids": [str(test_document)], "include_context": True, }, headers={"X-API-Key": test_key}, ) # Verify response assert response.status_code == 200 data = response.json() assert "answer" in data assert "sources" in data def test_query_without_context_returns_no_sources( self, test_client, test_api_key, mock_milvus_search, mock_azure_openai, mock_embeddings, ): """Test query with include_context=False returns no sources.""" test_key, _ = test_api_key # Send query without context response = test_client.post( "/query", json={ "query": "test query", "include_context": False, }, headers={"X-API-Key": test_key}, ) # Verify response assert response.status_code == 200 data = response.json() assert "answer" in data assert "sources" in data assert len(data["sources"]) == 0 def test_query_empty_string_returns_400(self, test_client, test_api_key): """Test empty query returns 400 validation error.""" test_key, _ = test_api_key # Send empty query response = test_client.post( "/query", json={"query": ""}, headers={"X-API-Key": test_key}, ) # Verify error response assert response.status_code == 422 # Pydantic validation error def test_query_whitespace_only_returns_400(self, test_client, test_api_key): """Test whitespace-only query returns 400 validation error.""" test_key, _ = test_api_key # Send whitespace query response = test_client.post( "/query", json={"query": " "}, headers={"X-API-Key": test_key}, ) # Verify error response assert response.status_code == 422 def test_query_too_long_returns_400(self, test_client, test_api_key): """Test query exceeding 1000 chars returns 400 validation error.""" test_key, _ = test_api_key # Send query exceeding max length long_query = "a" * 1001 response = test_client.post( "/query", json={"query": long_query}, headers={"X-API-Key": test_key}, ) # Verify error response assert response.status_code == 422 def test_query_missing_api_key_returns_401(self, test_client): """Test missing API key returns 401 unauthorized.""" # Send query without API key response = test_client.post( "/query", json={"query": "test query"}, ) # Verify error response assert response.status_code == 403 # API key dependency raises 403 def test_query_invalid_api_key_returns_401(self, test_client): """Test invalid API key returns 401 unauthorized.""" # Send query with invalid API key response = test_client.post( "/query", json={"query": "test query"}, headers={"X-API-Key": "invalid_key_12345"}, ) # Verify error response assert response.status_code == 403 def test_query_invalid_document_id_returns_404( self, test_client, test_api_key, mock_embeddings, ): """Test query with nonexistent document_id returns 404.""" test_key, _ = test_api_key # Send query with invalid document ID fake_doc_id = str(uuid4()) response = test_client.post( "/query", json={ "query": "test query", "document_ids": [fake_doc_id], }, headers={"X-API-Key": test_key}, ) # Verify error response assert response.status_code == 404 data = response.json() assert "detail" in data def test_query_processing_time_tracked( self, test_client, test_api_key, mock_milvus_search, mock_azure_openai, mock_embeddings, ): """Test processing time is tracked and included in response.""" test_key, _ = test_api_key # Send query response = test_client.post( "/query", json={"query": "test query"}, headers={"X-API-Key": test_key}, ) # Verify processing time assert response.status_code == 200 data = response.json() assert "processing_time_ms" in data assert isinstance(data["processing_time_ms"], int) assert data["processing_time_ms"] > 0 assert data["processing_time_ms"] < 10000 # Should be < 10 seconds def test_query_confidence_score_included( self, test_client, test_api_key, mock_milvus_search, mock_azure_openai, mock_embeddings, ): """Test confidence score is calculated and included.""" test_key, _ = test_api_key # Send query response = test_client.post( "/query", json={"query": "test query"}, headers={"X-API-Key": test_key}, ) # Verify confidence score assert response.status_code == 200 data = response.json() assert "confidence" in data assert isinstance(data["confidence"], float) assert 0.0 <= data["confidence"] <= 1.0 def test_query_no_results_returns_low_confidence( self, test_client, test_api_key, mock_embeddings, ): """Test query with no results returns low confidence answer.""" test_key, _ = test_api_key with patch("src.retrieval.semantic_retriever.connections"): with patch("src.retrieval.semantic_retriever.Collection") as mock_coll: # Mock empty search results mock_collection = MagicMock() mock_collection.load = MagicMock() mock_collection.search.return_value = [[]] mock_coll.return_value = mock_collection # Send query response = test_client.post( "/query", json={"query": "nonexistent query"}, headers={"X-API-Key": test_key}, ) # Verify low confidence response assert response.status_code == 200 data = response.json() assert "couldn't find any information" in data["answer"].lower() assert data["confidence"] < 0.2 def test_query_retrieval_error_returns_500( self, test_client, test_api_key, mock_embeddings, ): """Test retrieval error returns 500 internal error.""" test_key, _ = test_api_key with patch("src.retrieval.semantic_retriever.connections") as mock_conn: # Mock connection error mock_conn.connect.side_effect = Exception("Milvus connection failed") # Send query response = test_client.post( "/query", json={"query": "test query"}, headers={"X-API-Key": test_key}, ) # Verify error response assert response.status_code == 500 data = response.json() assert "detail" in data def test_query_generation_error_returns_500( self, test_client, test_api_key, mock_milvus_search, mock_embeddings, ): """Test answer generation error returns 500 internal error.""" test_key, _ = test_api_key with patch("src.retrieval.answer_generator.AsyncAzureOpenAI") as mock_class: mock_client = MagicMock() mock_client.chat.completions.create = AsyncMock( side_effect=Exception("OpenAI API error") ) mock_class.return_value = mock_client # Send query response = test_client.post( "/query", json={"query": "test query"}, headers={"X-API-Key": test_key}, ) # Verify error response assert response.status_code == 500 data = response.json() assert "detail" in data def test_query_strips_whitespace( self, test_client, test_api_key, mock_milvus_search, mock_azure_openai, mock_embeddings, ): """Test query whitespace is stripped.""" test_key, _ = test_api_key # Send query with leading/trailing whitespace response = test_client.post( "/query", json={"query": " test query "}, headers={"X-API-Key": test_key}, ) # Verify whitespace was stripped assert response.status_code == 200 data = response.json() assert data["query"] == "test query" def test_query_multiple_document_ids( self, test_client, test_api_key, test_document, mock_milvus_search, mock_azure_openai, mock_embeddings, ): """Test query with multiple document_ids.""" test_key, _ = test_api_key # Create second test document doc_id_2 = uuid4() with get_sync_session() as session: document_2 = Document( id=doc_id_2, filename="test_document_2.xlsx", status="completed", file_size=2048, mime_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", ) session.add(document_2) session.commit() try: # Send query with multiple document IDs response = test_client.post( "/query", json={ "query": "test query", "document_ids": [str(test_document), str(doc_id_2)], }, headers={"X-API-Key": test_key}, ) # Verify response assert response.status_code == 200 data = response.json() assert "answer" in data finally: # Cleanup second document with get_sync_session() as session: if doc_obj := session.execute( select(Document).where(Document.id == doc_id_2) ).scalar_one_or_none(): session.delete(doc_obj) session.commit()