Fixing tests

tests/integration/test_document_rag_integration.py

@@ -11,6 +11,14 @@ from unittest.mock import AsyncMock, MagicMock
 from trustgraph.retrieval.document_rag.document_rag import DocumentRag
 
 
+# Sample chunk content for testing - maps chunk_id to content
+CHUNK_CONTENT = {
+    "doc/c1": "Machine learning is a subset of artificial intelligence that focuses on algorithms that learn from data.",
+    "doc/c2": "Deep learning uses neural networks with multiple layers to model complex patterns in data.",
+    "doc/c3": "Supervised learning algorithms learn from labeled training data to make predictions on new data.",
+}
+
+
 @pytest.mark.integration
 class TestDocumentRagIntegration:
     """Integration tests for DocumentRAG system coordination"""
@@ -27,15 +35,19 @@ class TestDocumentRagIntegration:
 
     @pytest.fixture
     def mock_doc_embeddings_client(self):
-        """Mock document embeddings client that returns realistic document chunks"""
+        """Mock document embeddings client that returns chunk IDs"""
         client = AsyncMock()
-        client.query.return_value = [
+        # Now returns chunk_ids instead of actual content
-            "Machine learning is a subset of artificial intelligence that focuses on algorithms that learn from data.",
+        client.query.return_value = ["doc/c1", "doc/c2", "doc/c3"]
-            "Deep learning uses neural networks with multiple layers to model complex patterns in data.",
-            "Supervised learning algorithms learn from labeled training data to make predictions on new data."
-        ]
         return client
 
+    @pytest.fixture
+    def mock_fetch_chunk(self):
+        """Mock fetch_chunk function that retrieves chunk content from librarian"""
+        async def fetch(chunk_id, user):
+            return CHUNK_CONTENT.get(chunk_id, f"Content for {chunk_id}")
+        return fetch
+
     @pytest.fixture
     def mock_prompt_client(self):
         """Mock prompt client that generates realistic responses"""
@@ -48,12 +60,14 @@ class TestDocumentRagIntegration:
         return client
 
     @pytest.fixture
-    def document_rag(self, mock_embeddings_client, mock_doc_embeddings_client, mock_prompt_client):
+    def document_rag(self, mock_embeddings_client, mock_doc_embeddings_client,
+                     mock_prompt_client, mock_fetch_chunk):
         """Create DocumentRag instance with mocked dependencies"""
         return DocumentRag(
             embeddings_client=mock_embeddings_client,
             doc_embeddings_client=mock_doc_embeddings_client,
             prompt_client=mock_prompt_client,
+            fetch_chunk=mock_fetch_chunk,
             verbose=True
         )
 
@@ -85,6 +99,7 @@ class TestDocumentRagIntegration:
             collection=collection
         )
 
+        # Documents are fetched from librarian using chunk_ids
         mock_prompt_client.document_prompt.assert_called_once_with(
             query=query,
             documents=[
@@ -102,16 +117,18 @@ class TestDocumentRagIntegration:
 
     @pytest.mark.asyncio
     async def test_document_rag_with_no_documents_found(self, mock_embeddings_client,
-                                                        mock_doc_embeddings_client, mock_prompt_client):
+                                                        mock_doc_embeddings_client, mock_prompt_client,
+                                                        mock_fetch_chunk):
         """Test DocumentRAG behavior when no documents are retrieved"""
         # Arrange
-        mock_doc_embeddings_client.query.return_value = []  # No documents found
+        mock_doc_embeddings_client.query.return_value = []  # No chunk_ids found
         mock_prompt_client.document_prompt.return_value = "I couldn't find any relevant documents for your query."
 
         document_rag = DocumentRag(
             embeddings_client=mock_embeddings_client,
             doc_embeddings_client=mock_doc_embeddings_client,
             prompt_client=mock_prompt_client,
+            fetch_chunk=mock_fetch_chunk,
             verbose=False
         )
 
@@ -130,7 +147,8 @@ class TestDocumentRagIntegration:
 
     @pytest.mark.asyncio
     async def test_document_rag_embeddings_service_failure(self, mock_embeddings_client,
-                                                          mock_doc_embeddings_client, mock_prompt_client):
+                                                          mock_doc_embeddings_client, mock_prompt_client,
+                                                          mock_fetch_chunk):
         """Test DocumentRAG error handling when embeddings service fails"""
         # Arrange
         mock_embeddings_client.embed.side_effect = Exception("Embeddings service unavailable")
@@ -139,6 +157,7 @@ class TestDocumentRagIntegration:
             embeddings_client=mock_embeddings_client,
             doc_embeddings_client=mock_doc_embeddings_client,
             prompt_client=mock_prompt_client,
+            fetch_chunk=mock_fetch_chunk,
             verbose=False
         )
 
@@ -153,7 +172,8 @@ class TestDocumentRagIntegration:
 
     @pytest.mark.asyncio
     async def test_document_rag_document_service_failure(self, mock_embeddings_client,
-                                                        mock_doc_embeddings_client, mock_prompt_client):
+                                                        mock_doc_embeddings_client, mock_prompt_client,
+                                                        mock_fetch_chunk):
         """Test DocumentRAG error handling when document service fails"""
         # Arrange
         mock_doc_embeddings_client.query.side_effect = Exception("Document service connection failed")
@@ -162,6 +182,7 @@ class TestDocumentRagIntegration:
             embeddings_client=mock_embeddings_client,
             doc_embeddings_client=mock_doc_embeddings_client,
             prompt_client=mock_prompt_client,
+            fetch_chunk=mock_fetch_chunk,
             verbose=False
         )
 
@@ -176,7 +197,8 @@ class TestDocumentRagIntegration:
 
     @pytest.mark.asyncio
     async def test_document_rag_prompt_service_failure(self, mock_embeddings_client,
-                                                      mock_doc_embeddings_client, mock_prompt_client):
+                                                      mock_doc_embeddings_client, mock_prompt_client,
+                                                      mock_fetch_chunk):
         """Test DocumentRAG error handling when prompt service fails"""
         # Arrange
         mock_prompt_client.document_prompt.side_effect = Exception("LLM service rate limited")
@@ -185,6 +207,7 @@ class TestDocumentRagIntegration:
             embeddings_client=mock_embeddings_client,
             doc_embeddings_client=mock_doc_embeddings_client,
             prompt_client=mock_prompt_client,
+            fetch_chunk=mock_fetch_chunk,
             verbose=False
         )
 
@@ -247,6 +270,7 @@ class TestDocumentRagIntegration:
     @pytest.mark.asyncio
     async def test_document_rag_verbose_logging(self, mock_embeddings_client,
                                                mock_doc_embeddings_client, mock_prompt_client,
+                                               mock_fetch_chunk,
                                                caplog):
         """Test DocumentRAG verbose logging functionality"""
         import logging
@@ -258,6 +282,7 @@ class TestDocumentRagIntegration:
             embeddings_client=mock_embeddings_client,
             doc_embeddings_client=mock_doc_embeddings_client,
             prompt_client=mock_prompt_client,
+            fetch_chunk=mock_fetch_chunk,
             verbose=True
         )
 
@@ -269,7 +294,7 @@ class TestDocumentRagIntegration:
         assert "DocumentRag initialized" in log_messages
         assert "Constructing prompt..." in log_messages
         assert "Computing embeddings..." in log_messages
-        assert "Getting documents..." in log_messages
+        assert "chunk_ids" in log_messages.lower()
         assert "Invoking LLM..." in log_messages
         assert "Query processing complete" in log_messages
 
@@ -278,9 +303,9 @@ class TestDocumentRagIntegration:
     async def test_document_rag_performance_with_large_document_set(self, document_rag,
                                                                    mock_doc_embeddings_client):
         """Test DocumentRAG performance with large document retrieval"""
-        # Arrange - Mock large document set (100 documents)
+        # Arrange - Mock large chunk_id set (100 chunks)
-        large_doc_set = [f"Document {i} content about machine learning and AI" for i in range(100)]
+        large_chunk_ids = [f"doc/c{i}" for i in range(100)]
-        mock_doc_embeddings_client.query.return_value = large_doc_set
+        mock_doc_embeddings_client.query.return_value = large_chunk_ids
 
         # Act
         import time

tests/integration/test_document_rag_streaming_integration.py

@@ -14,6 +14,14 @@ from tests.utils.streaming_assertions import (
 )
 
 
+# Sample chunk content for testing - maps chunk_id to content
+CHUNK_CONTENT = {
+    "doc/c1": "Machine learning is a subset of AI.",
+    "doc/c2": "Deep learning uses neural networks.",
+    "doc/c3": "Supervised learning needs labeled data.",
+}
+
+
 @pytest.mark.integration
 class TestDocumentRagStreaming:
     """Integration tests for DocumentRAG streaming"""
@@ -27,15 +35,19 @@ class TestDocumentRagStreaming:
 
     @pytest.fixture
     def mock_doc_embeddings_client(self):
-        """Mock document embeddings client"""
+        """Mock document embeddings client that returns chunk IDs"""
         client = AsyncMock()
-        client.query.return_value = [
+        # Now returns chunk_ids instead of actual content
-            "Machine learning is a subset of AI.",
+        client.query.return_value = ["doc/c1", "doc/c2", "doc/c3"]
-            "Deep learning uses neural networks.",
-            "Supervised learning needs labeled data."
-        ]
         return client
 
+    @pytest.fixture
+    def mock_fetch_chunk(self):
+        """Mock fetch_chunk function that retrieves chunk content from librarian"""
+        async def fetch(chunk_id, user):
+            return CHUNK_CONTENT.get(chunk_id, f"Content for {chunk_id}")
+        return fetch
+
     @pytest.fixture
     def mock_streaming_prompt_client(self, mock_streaming_llm_response):
         """Mock prompt client with streaming support"""
@@ -66,12 +78,13 @@ class TestDocumentRagStreaming:
 
     @pytest.fixture
     def document_rag_streaming(self, mock_embeddings_client, mock_doc_embeddings_client,
-                                mock_streaming_prompt_client):
+                                mock_streaming_prompt_client, mock_fetch_chunk):
         """Create DocumentRag instance with streaming support"""
         return DocumentRag(
             embeddings_client=mock_embeddings_client,
             doc_embeddings_client=mock_doc_embeddings_client,
             prompt_client=mock_streaming_prompt_client,
+            fetch_chunk=mock_fetch_chunk,
             verbose=True
         )
 
@@ -190,7 +203,7 @@ class TestDocumentRagStreaming:
                                                              mock_doc_embeddings_client):
         """Test streaming with no documents found"""
         # Arrange
-        mock_doc_embeddings_client.query.return_value = []  # No documents
+        mock_doc_embeddings_client.query.return_value = []  # No chunk_ids
         callback = AsyncMock()
 
         # Act

tests/integration/test_rag_streaming_protocol.py

@@ -202,11 +202,18 @@ class TestDocumentRagStreamingProtocol:
 
     @pytest.fixture
     def mock_doc_embeddings_client(self):
-        """Mock document embeddings client"""
+        """Mock document embeddings client that returns chunk IDs"""
         client = AsyncMock()
-        client.query.return_value = ["doc1", "doc2"]
+        client.query.return_value = ["doc/c1", "doc/c2"]
         return client
 
+    @pytest.fixture
+    def mock_fetch_chunk(self):
+        """Mock fetch_chunk function that retrieves chunk content from librarian"""
+        async def fetch(chunk_id, user):
+            return f"Content for {chunk_id}"
+        return fetch
+
     @pytest.fixture
     def mock_streaming_prompt_client(self):
         """Mock prompt client with streaming support"""