"""Integration tests for DomainAwareRetriever. Note: These tests require a running Milvus instance and are skipped if Milvus is not available. """ import pytest from unittest.mock import AsyncMock, MagicMock, patch from src.knowledge.domain_aware_retriever import DomainAwareRetriever from src.knowledge.query_term_matcher import QueryTermMatcher @pytest.fixture def mock_query_matcher(): """Create a mock QueryTermMatcher.""" matcher = AsyncMock(spec=QueryTermMatcher) matcher.extract_query_terms = AsyncMock(return_value={ "matched_terms": ["revenue", "ebitda"], "domain_category": "finance", "matching_method": "keyword" }) return matcher @pytest.fixture def mock_embedding_service(): """Create a mock EmbeddingService.""" service = AsyncMock() # Return a dummy 3072-dim vector service.embed = AsyncMock(return_value=[0.1] * 3072) return service @pytest.fixture def retriever(mock_query_matcher, mock_embedding_service): """Create DomainAwareRetriever with mocked dependencies.""" return DomainAwareRetriever( query_matcher=mock_query_matcher, embedding_service=mock_embedding_service, ) class TestDomainAwareRetriever: """Unit tests for DomainAwareRetriever.""" def test_build_filter_with_terms_and_category(self, retriever): """Should build filter with category and terms.""" matched_terms = ["revenue", "ebitda", "profit"] category = "finance" filter_expr = retriever._build_filter( matched_terms, category, "tenant1", None ) assert filter_expr is not None assert "domain_category ==" in filter_expr assert "finance" in filter_expr assert "array_contains_any" in filter_expr assert "revenue" in filter_expr def test_build_filter_with_document_id(self, retriever): """Should include document_id in filter when provided.""" doc_id = "123e4567-e89b-12d3-a456-426614174000" filter_expr = retriever._build_filter( ["revenue"], "finance", "tenant1", doc_id ) assert filter_expr is not None assert f'document_id == "{doc_id}"' in filter_expr def test_build_filter_general_category_skipped(self, retriever): """Should not filter by category when category is 'general'.""" filter_expr = retriever._build_filter( ["term1"], "general", "tenant1", None ) # Should only have terms filter, not category filter if filter_expr: assert "domain_category" not in filter_expr assert "array_contains_any" in filter_expr def test_build_filter_no_terms_or_category(self, retriever): """Should return None when no terms or category.""" filter_expr = retriever._build_filter( [], "general", "tenant1", None ) assert filter_expr is None def test_build_filter_escapes_quotes(self, retriever): """Should escape quotes in terms.""" matched_terms = ["term's", 'term"s'] filter_expr = retriever._build_filter( matched_terms, "finance", "tenant1", None ) # Should have escaped the quotes assert filter_expr is not None # The single quote should be escaped assert "term" in filter_expr def test_build_filter_only_document_id(self, retriever): """Should build filter with only document_id.""" doc_id = "123e4567-e89b-12d3-a456-426614174000" filter_expr = retriever._build_filter( [], "general", "tenant1", doc_id ) assert filter_expr == f'document_id == "{doc_id}"' @pytest.mark.asyncio @patch("src.knowledge.domain_aware_retriever.connections") @patch("src.knowledge.domain_aware_retriever.Collection") async def test_retrieve_uses_filter( self, mock_collection_class, mock_connections, mock_query_matcher, mock_embedding_service ): """Should pass filter expression to Milvus search.""" # Mock Milvus collection mock_collection = MagicMock() mock_search_results = [[]] # Empty results mock_collection.search.return_value = mock_search_results mock_collection_class.return_value = mock_collection retriever = DomainAwareRetriever( query_matcher=mock_query_matcher, embedding_service=mock_embedding_service, ) await retriever.retrieve("What is Q4 EBITDA?") # Verify search was called with a filter assert mock_collection.search.called call_kwargs = mock_collection.search.call_args.kwargs assert "expr" in call_kwargs # Should have a filter (not None) since we have matched terms assert call_kwargs["expr"] is not None @pytest.mark.asyncio @patch("src.knowledge.domain_aware_retriever.connections") @patch("src.knowledge.domain_aware_retriever.Collection") async def test_retrieve_formats_results( self, mock_collection_class, mock_connections, mock_query_matcher, mock_embedding_service ): """Should format Milvus results correctly.""" # Mock Milvus search results mock_hit = MagicMock() mock_hit.entity.get.side_effect = lambda key: { "record_id": "rec_123", "document_id": "doc_456", "filename": "test.xlsx", "document_type": "excel", "text_content": "Revenue was $1M", "metadata": {"sheet_name": "Sheet1"}, "domain_category": "finance", "domain_terms": ["revenue"], }.get(key) mock_hit.score = 0.95 mock_collection = MagicMock() mock_collection.search.return_value = [[mock_hit]] mock_collection_class.return_value = mock_collection retriever = DomainAwareRetriever( query_matcher=mock_query_matcher, embedding_service=mock_embedding_service, ) results = await retriever.retrieve("What is revenue?") assert len(results) == 1 assert results[0]["record_id"] == "rec_123" assert results[0]["document_id"] == "doc_456" assert results[0]["text_content"] == "Revenue was $1M" assert results[0]["domain_category"] == "finance" assert results[0]["score"] == 0.95 @pytest.mark.asyncio async def test_retrieve_with_no_matches( self, mock_query_matcher, mock_embedding_service ): """Should handle case when query matcher returns no terms.""" # Configure matcher to return no terms mock_query_matcher.extract_query_terms = AsyncMock(return_value={ "matched_terms": [], "domain_category": "general", "matching_method": "none" }) with patch("src.knowledge.domain_aware_retriever.connections"): with patch("src.knowledge.domain_aware_retriever.Collection") as mock_coll: mock_collection = MagicMock() mock_collection.search.return_value = [[]] mock_coll.return_value = mock_collection retriever = DomainAwareRetriever( query_matcher=mock_query_matcher, embedding_service=mock_embedding_service, ) results = await retriever.retrieve("generic query") # Should still perform search, but without filter assert mock_collection.search.called call_kwargs = mock_collection.search.call_args.kwargs assert call_kwargs["expr"] is None # No filter @pytest.mark.asyncio @patch("src.knowledge.domain_aware_retriever.connections") @patch("src.knowledge.domain_aware_retriever.Collection") async def test_retrieve_respects_top_k( self, mock_collection_class, mock_connections, mock_query_matcher, mock_embedding_service ): """Should pass top_k parameter to Milvus search.""" mock_collection = MagicMock() mock_collection.search.return_value = [[]] mock_collection_class.return_value = mock_collection retriever = DomainAwareRetriever( query_matcher=mock_query_matcher, embedding_service=mock_embedding_service, ) await retriever.retrieve("test query", top_k=25) # Verify limit parameter call_kwargs = mock_collection.search.call_args.kwargs assert call_kwargs["limit"] == 25 class TestIntegration: """Integration tests (require running Milvus instance).""" @pytest.mark.integration @pytest.mark.asyncio async def test_full_retrieval_pipeline(self): """Test complete retrieval flow with real Milvus. Note: This test is marked as 'integration' and requires: - Running Milvus instance - Populated collection with domain-aware data """ # This test would require actual Milvus setup # Skip for now - would be run in integration test suite pytest.skip("Requires running Milvus instance with test data") @pytest.mark.integration @pytest.mark.asyncio async def test_performance_with_filtering_vs_full_search(self): """Benchmark filtered vs full search performance. Note: This test is marked as 'integration' and requires: - Running Milvus instance - Large dataset (100k+ chunks) """ pytest.skip("Requires running Milvus instance with large dataset")