From 35dac42c73a224fd90172649fee95b5c2bac5f80 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Sat, 7 Mar 2026 23:36:02 +0000 Subject: [PATCH] Fixing tests --- .../unit/test_retrieval/test_document_rag.py | 280 +++++++++++------- .../test_doc_embeddings_milvus_storage.py | 39 ++- 2 files changed, 192 insertions(+), 127 deletions(-) diff --git a/tests/unit/test_retrieval/test_document_rag.py b/tests/unit/test_retrieval/test_document_rag.py index 590572bc..c9f5e8e1 100644 --- a/tests/unit/test_retrieval/test_document_rag.py +++ b/tests/unit/test_retrieval/test_document_rag.py @@ -8,48 +8,75 @@ from unittest.mock import MagicMock, AsyncMock from trustgraph.retrieval.document_rag.document_rag import DocumentRag, Query +# Sample chunk content mapping for tests +CHUNK_CONTENT = { + "doc/c1": "Document 1 content", + "doc/c2": "Document 2 content", + "doc/c3": "Relevant document content", + "doc/c4": "Another document", + "doc/c5": "Default doc", + "doc/c6": "Verbose test doc", + "doc/c7": "Verbose doc content", + "doc/ml1": "Machine learning is a subset of artificial intelligence...", + "doc/ml2": "ML algorithms learn patterns from data to make predictions...", + "doc/ml3": "Common ML techniques include supervised and unsupervised learning...", +} + + +@pytest.fixture +def mock_fetch_chunk(): + """Create a mock fetch_chunk function""" + async def fetch(chunk_id, user): + return CHUNK_CONTENT.get(chunk_id, f"Content for {chunk_id}") + return fetch + + class TestDocumentRag: """Test cases for DocumentRag class""" - def test_document_rag_initialization_with_defaults(self): + def test_document_rag_initialization_with_defaults(self, mock_fetch_chunk): """Test DocumentRag initialization with default verbose setting""" # Create mock clients mock_prompt_client = MagicMock() mock_embeddings_client = MagicMock() mock_doc_embeddings_client = MagicMock() - + # Initialize DocumentRag document_rag = DocumentRag( prompt_client=mock_prompt_client, embeddings_client=mock_embeddings_client, - doc_embeddings_client=mock_doc_embeddings_client + doc_embeddings_client=mock_doc_embeddings_client, + fetch_chunk=mock_fetch_chunk ) - + # Verify initialization assert document_rag.prompt_client == mock_prompt_client assert document_rag.embeddings_client == mock_embeddings_client assert document_rag.doc_embeddings_client == mock_doc_embeddings_client + assert document_rag.fetch_chunk == mock_fetch_chunk assert document_rag.verbose is False # Default value - def test_document_rag_initialization_with_verbose(self): + def test_document_rag_initialization_with_verbose(self, mock_fetch_chunk): """Test DocumentRag initialization with verbose enabled""" # Create mock clients mock_prompt_client = MagicMock() mock_embeddings_client = MagicMock() mock_doc_embeddings_client = MagicMock() - + # Initialize DocumentRag with verbose=True document_rag = DocumentRag( prompt_client=mock_prompt_client, embeddings_client=mock_embeddings_client, doc_embeddings_client=mock_doc_embeddings_client, + fetch_chunk=mock_fetch_chunk, verbose=True ) - + # Verify initialization assert document_rag.prompt_client == mock_prompt_client assert document_rag.embeddings_client == mock_embeddings_client assert document_rag.doc_embeddings_client == mock_doc_embeddings_client + assert document_rag.fetch_chunk == mock_fetch_chunk assert document_rag.verbose is True @@ -60,7 +87,7 @@ class TestQuery: """Test Query initialization with default parameters""" # Create mock DocumentRag mock_rag = MagicMock() - + # Initialize Query with defaults query = Query( rag=mock_rag, @@ -68,7 +95,7 @@ class TestQuery: collection="test_collection", verbose=False ) - + # Verify initialization assert query.rag == mock_rag assert query.user == "test_user" @@ -80,7 +107,7 @@ class TestQuery: """Test Query initialization with custom doc_limit""" # Create mock DocumentRag mock_rag = MagicMock() - + # Initialize Query with custom doc_limit query = Query( rag=mock_rag, @@ -89,7 +116,7 @@ class TestQuery: verbose=True, doc_limit=50 ) - + # Verify initialization assert query.rag == mock_rag assert query.user == "custom_user" @@ -104,11 +131,11 @@ class TestQuery: mock_rag = MagicMock() mock_embeddings_client = AsyncMock() mock_rag.embeddings_client = mock_embeddings_client - + # Mock the embed method to return test vectors expected_vectors = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]] mock_embeddings_client.embed.return_value = expected_vectors - + # Initialize Query query = Query( rag=mock_rag, @@ -116,14 +143,14 @@ class TestQuery: collection="test_collection", verbose=False ) - + # Call get_vector test_query = "What documents are relevant?" result = await query.get_vector(test_query) - + # Verify embeddings client was called correctly mock_embeddings_client.embed.assert_called_once_with(test_query) - + # Verify result matches expected vectors assert result == expected_vectors @@ -136,15 +163,20 @@ class TestQuery: mock_doc_embeddings_client = AsyncMock() mock_rag.embeddings_client = mock_embeddings_client mock_rag.doc_embeddings_client = mock_doc_embeddings_client - + + # Mock fetch_chunk function + async def mock_fetch(chunk_id, user): + return CHUNK_CONTENT.get(chunk_id, f"Content for {chunk_id}") + mock_rag.fetch_chunk = mock_fetch + # Mock the embedding and document query responses test_vectors = [[0.1, 0.2, 0.3]] mock_embeddings_client.embed.return_value = test_vectors - - # Mock document results - test_docs = ["Document 1 content", "Document 2 content"] - mock_doc_embeddings_client.query.return_value = test_docs - + + # Mock document embeddings returns chunk_ids + test_chunk_ids = ["doc/c1", "doc/c2"] + mock_doc_embeddings_client.query.return_value = test_chunk_ids + # Initialize Query query = Query( rag=mock_rag, @@ -153,14 +185,14 @@ class TestQuery: verbose=False, doc_limit=15 ) - + # Call get_docs test_query = "Find relevant documents" result = await query.get_docs(test_query) - + # Verify embeddings client was called mock_embeddings_client.embed.assert_called_once_with(test_query) - + # Verify doc embeddings client was called correctly mock_doc_embeddings_client.query.assert_called_once_with( test_vectors, @@ -168,35 +200,37 @@ class TestQuery: user="test_user", collection="test_collection" ) - - # Verify result is list of documents - assert result == test_docs + + # Verify result is list of fetched document content + assert "Document 1 content" in result + assert "Document 2 content" in result @pytest.mark.asyncio - async def test_document_rag_query_method(self): + async def test_document_rag_query_method(self, mock_fetch_chunk): """Test DocumentRag.query method orchestrates full document RAG pipeline""" # Create mock clients mock_prompt_client = AsyncMock() mock_embeddings_client = AsyncMock() mock_doc_embeddings_client = AsyncMock() - - # Mock embeddings and document responses + + # Mock embeddings and document embeddings responses test_vectors = [[0.1, 0.2, 0.3]] - test_docs = ["Relevant document content", "Another document"] + test_chunk_ids = ["doc/c3", "doc/c4"] expected_response = "This is the document RAG response" - + mock_embeddings_client.embed.return_value = test_vectors - mock_doc_embeddings_client.query.return_value = test_docs + mock_doc_embeddings_client.query.return_value = test_chunk_ids mock_prompt_client.document_prompt.return_value = expected_response - + # Initialize DocumentRag document_rag = DocumentRag( prompt_client=mock_prompt_client, embeddings_client=mock_embeddings_client, doc_embeddings_client=mock_doc_embeddings_client, + fetch_chunk=mock_fetch_chunk, verbose=False ) - + # Call DocumentRag.query result = await document_rag.query( query="test query", @@ -204,10 +238,10 @@ class TestQuery: collection="test_collection", doc_limit=10 ) - + # Verify embeddings client was called mock_embeddings_client.embed.assert_called_once_with("test query") - + # Verify doc embeddings client was called mock_doc_embeddings_client.query.assert_called_once_with( test_vectors, @@ -215,39 +249,43 @@ class TestQuery: user="test_user", collection="test_collection" ) - - # Verify prompt client was called with documents and query - mock_prompt_client.document_prompt.assert_called_once_with( - query="test query", - documents=test_docs - ) - + + # Verify prompt client was called with fetched documents and query + mock_prompt_client.document_prompt.assert_called_once() + call_args = mock_prompt_client.document_prompt.call_args + assert call_args.kwargs["query"] == "test query" + # Documents should be fetched content, not chunk_ids + docs = call_args.kwargs["documents"] + assert "Relevant document content" in docs + assert "Another document" in docs + # Verify result assert result == expected_response @pytest.mark.asyncio - async def test_document_rag_query_with_defaults(self): + async def test_document_rag_query_with_defaults(self, mock_fetch_chunk): """Test DocumentRag.query method with default parameters""" # Create mock clients mock_prompt_client = AsyncMock() mock_embeddings_client = AsyncMock() mock_doc_embeddings_client = AsyncMock() - + # Mock responses mock_embeddings_client.embed.return_value = [[0.1, 0.2]] - mock_doc_embeddings_client.query.return_value = ["Default doc"] + mock_doc_embeddings_client.query.return_value = ["doc/c5"] mock_prompt_client.document_prompt.return_value = "Default response" - + # Initialize DocumentRag document_rag = DocumentRag( prompt_client=mock_prompt_client, embeddings_client=mock_embeddings_client, - doc_embeddings_client=mock_doc_embeddings_client + doc_embeddings_client=mock_doc_embeddings_client, + fetch_chunk=mock_fetch_chunk ) - + # Call DocumentRag.query with minimal parameters result = await document_rag.query("simple query") - + # Verify default parameters were used mock_doc_embeddings_client.query.assert_called_once_with( [[0.1, 0.2]], @@ -255,7 +293,7 @@ class TestQuery: user="trustgraph", # Default user collection="default" # Default collection ) - + assert result == "Default response" @pytest.mark.asyncio @@ -267,11 +305,16 @@ class TestQuery: mock_doc_embeddings_client = AsyncMock() mock_rag.embeddings_client = mock_embeddings_client mock_rag.doc_embeddings_client = mock_doc_embeddings_client - + + # Mock fetch_chunk + async def mock_fetch(chunk_id, user): + return CHUNK_CONTENT.get(chunk_id, f"Content for {chunk_id}") + mock_rag.fetch_chunk = mock_fetch + # Mock responses mock_embeddings_client.embed.return_value = [[0.7, 0.8]] - mock_doc_embeddings_client.query.return_value = ["Verbose test doc"] - + mock_doc_embeddings_client.query.return_value = ["doc/c6"] + # Initialize Query with verbose=True query = Query( rag=mock_rag, @@ -280,49 +323,51 @@ class TestQuery: verbose=True, doc_limit=5 ) - + # Call get_docs result = await query.get_docs("verbose test") - + # Verify calls were made mock_embeddings_client.embed.assert_called_once_with("verbose test") mock_doc_embeddings_client.query.assert_called_once() - - # Verify result - assert result == ["Verbose test doc"] + + # Verify result contains fetched content + assert "Verbose test doc" in result @pytest.mark.asyncio - async def test_document_rag_query_with_verbose(self): + async def test_document_rag_query_with_verbose(self, mock_fetch_chunk): """Test DocumentRag.query method with verbose logging enabled""" # Create mock clients mock_prompt_client = AsyncMock() mock_embeddings_client = AsyncMock() mock_doc_embeddings_client = AsyncMock() - + # Mock responses mock_embeddings_client.embed.return_value = [[0.3, 0.4]] - mock_doc_embeddings_client.query.return_value = ["Verbose doc content"] + mock_doc_embeddings_client.query.return_value = ["doc/c7"] mock_prompt_client.document_prompt.return_value = "Verbose RAG response" - + # Initialize DocumentRag with verbose=True document_rag = DocumentRag( prompt_client=mock_prompt_client, embeddings_client=mock_embeddings_client, doc_embeddings_client=mock_doc_embeddings_client, + fetch_chunk=mock_fetch_chunk, verbose=True ) - + # Call DocumentRag.query result = await document_rag.query("verbose query test") - + # Verify all clients were called mock_embeddings_client.embed.assert_called_once_with("verbose query test") mock_doc_embeddings_client.query.assert_called_once() - mock_prompt_client.document_prompt.assert_called_once_with( - query="verbose query test", - documents=["Verbose doc content"] - ) - + + # Verify prompt client was called with fetched content + call_args = mock_prompt_client.document_prompt.call_args + assert call_args.kwargs["query"] == "verbose query test" + assert "Verbose doc content" in call_args.kwargs["documents"] + assert result == "Verbose RAG response" @pytest.mark.asyncio @@ -334,11 +379,16 @@ class TestQuery: mock_doc_embeddings_client = AsyncMock() mock_rag.embeddings_client = mock_embeddings_client mock_rag.doc_embeddings_client = mock_doc_embeddings_client - - # Mock responses - empty document list + + # Mock fetch_chunk (won't be called if no chunk_ids) + async def mock_fetch(chunk_id, user): + return f"Content for {chunk_id}" + mock_rag.fetch_chunk = mock_fetch + + # Mock responses - empty chunk_id list mock_embeddings_client.embed.return_value = [[0.1, 0.2]] - mock_doc_embeddings_client.query.return_value = [] # No documents found - + mock_doc_embeddings_client.query.return_value = [] # No chunk_ids found + # Initialize Query query = Query( rag=mock_rag, @@ -346,47 +396,48 @@ class TestQuery: collection="test_collection", verbose=False ) - + # Call get_docs result = await query.get_docs("query with no results") - + # Verify calls were made mock_embeddings_client.embed.assert_called_once_with("query with no results") mock_doc_embeddings_client.query.assert_called_once() - + # Verify empty result is returned assert result == [] @pytest.mark.asyncio - async def test_document_rag_query_with_empty_documents(self): + async def test_document_rag_query_with_empty_documents(self, mock_fetch_chunk): """Test DocumentRag.query method when no documents are retrieved""" # Create mock clients mock_prompt_client = AsyncMock() mock_embeddings_client = AsyncMock() mock_doc_embeddings_client = AsyncMock() - - # Mock responses - no documents found + + # Mock responses - no chunk_ids found mock_embeddings_client.embed.return_value = [[0.5, 0.6]] - mock_doc_embeddings_client.query.return_value = [] # Empty document list + mock_doc_embeddings_client.query.return_value = [] # Empty chunk_id list mock_prompt_client.document_prompt.return_value = "No documents found response" - + # Initialize DocumentRag document_rag = DocumentRag( prompt_client=mock_prompt_client, embeddings_client=mock_embeddings_client, doc_embeddings_client=mock_doc_embeddings_client, + fetch_chunk=mock_fetch_chunk, verbose=False ) - + # Call DocumentRag.query result = await document_rag.query("query with no matching docs") - + # Verify prompt client was called with empty document list mock_prompt_client.document_prompt.assert_called_once_with( query="query with no matching docs", documents=[] ) - + assert result == "No documents found response" @pytest.mark.asyncio @@ -396,11 +447,11 @@ class TestQuery: mock_rag = MagicMock() mock_embeddings_client = AsyncMock() mock_rag.embeddings_client = mock_embeddings_client - + # Mock the embed method expected_vectors = [[0.9, 1.0, 1.1]] mock_embeddings_client.embed.return_value = expected_vectors - + # Initialize Query with verbose=True query = Query( rag=mock_rag, @@ -408,68 +459,71 @@ class TestQuery: collection="test_collection", verbose=True ) - + # Call get_vector result = await query.get_vector("verbose vector test") - + # Verify embeddings client was called mock_embeddings_client.embed.assert_called_once_with("verbose vector test") - + # Verify result assert result == expected_vectors @pytest.mark.asyncio - async def test_document_rag_integration_flow(self): + async def test_document_rag_integration_flow(self, mock_fetch_chunk): """Test complete DocumentRag integration with realistic data flow""" # Create mock clients mock_prompt_client = AsyncMock() mock_embeddings_client = AsyncMock() mock_doc_embeddings_client = AsyncMock() - + # Mock realistic responses query_text = "What is machine learning?" query_vectors = [[0.1, 0.2, 0.3, 0.4, 0.5]] - retrieved_docs = [ - "Machine learning is a subset of artificial intelligence...", - "ML algorithms learn patterns from data to make predictions...", - "Common ML techniques include supervised and unsupervised learning..." - ] + retrieved_chunk_ids = ["doc/ml1", "doc/ml2", "doc/ml3"] final_response = "Machine learning is a field of AI that enables computers to learn and improve from experience without being explicitly programmed." - + mock_embeddings_client.embed.return_value = query_vectors - mock_doc_embeddings_client.query.return_value = retrieved_docs + mock_doc_embeddings_client.query.return_value = retrieved_chunk_ids mock_prompt_client.document_prompt.return_value = final_response - + # Initialize DocumentRag document_rag = DocumentRag( prompt_client=mock_prompt_client, embeddings_client=mock_embeddings_client, doc_embeddings_client=mock_doc_embeddings_client, + fetch_chunk=mock_fetch_chunk, verbose=False ) - + # Execute full pipeline result = await document_rag.query( query=query_text, - user="research_user", + user="research_user", collection="ml_knowledge", doc_limit=25 ) - + # Verify complete pipeline execution mock_embeddings_client.embed.assert_called_once_with(query_text) - + mock_doc_embeddings_client.query.assert_called_once_with( query_vectors, limit=25, user="research_user", collection="ml_knowledge" ) - - mock_prompt_client.document_prompt.assert_called_once_with( - query=query_text, - documents=retrieved_docs - ) - + + # Verify prompt client was called with fetched document content + mock_prompt_client.document_prompt.assert_called_once() + call_args = mock_prompt_client.document_prompt.call_args + assert call_args.kwargs["query"] == query_text + + # Verify documents were fetched from chunk_ids + docs = call_args.kwargs["documents"] + assert "Machine learning is a subset of artificial intelligence..." in docs + assert "ML algorithms learn patterns from data to make predictions..." in docs + assert "Common ML techniques include supervised and unsupervised learning..." in docs + # Verify final result - assert result == final_response \ No newline at end of file + assert result == final_response diff --git a/tests/unit/test_storage/test_doc_embeddings_milvus_storage.py b/tests/unit/test_storage/test_doc_embeddings_milvus_storage.py index 8da88fcb..ef66b741 100644 --- a/tests/unit/test_storage/test_doc_embeddings_milvus_storage.py +++ b/tests/unit/test_storage/test_doc_embeddings_milvus_storage.py @@ -148,7 +148,7 @@ class TestMilvusDocEmbeddingsStorageProcessor: @pytest.mark.asyncio async def test_store_document_embeddings_none_chunk(self, processor): - """Test storing document embeddings with None chunk_id (should be skipped)""" + """Test storing document embeddings with None chunk_id""" message = MagicMock() message.metadata = MagicMock() message.metadata.user = 'test_user' @@ -162,17 +162,19 @@ class TestMilvusDocEmbeddingsStorageProcessor: await processor.store_document_embeddings(message) - # Verify no insert was called for None chunk_id - processor.vecstore.insert.assert_not_called() + # Note: Implementation passes through None chunk_ids (only skips empty string "") + processor.vecstore.insert.assert_called_once_with( + [0.1, 0.2, 0.3], None, 'test_user', 'test_collection' + ) @pytest.mark.asyncio async def test_store_document_embeddings_mixed_valid_invalid_chunks(self, processor): - """Test storing document embeddings with mix of valid and invalid chunks""" + """Test storing document embeddings with mix of valid and empty chunks""" message = MagicMock() message.metadata = MagicMock() message.metadata.user = 'test_user' message.metadata.collection = 'test_collection' - + valid_chunk = ChunkEmbeddings( chunk_id="Valid document content", vectors=[[0.1, 0.2, 0.3]] @@ -181,18 +183,27 @@ class TestMilvusDocEmbeddingsStorageProcessor: chunk_id="", vectors=[[0.4, 0.5, 0.6]] ) - none_chunk = ChunkEmbeddings( - chunk_id=None, + another_valid = ChunkEmbeddings( + chunk_id="Another valid chunk", vectors=[[0.7, 0.8, 0.9]] ) - message.chunks = [valid_chunk, empty_chunk, none_chunk] - + message.chunks = [valid_chunk, empty_chunk, another_valid] + await processor.store_document_embeddings(message) - - # Verify only valid chunk was inserted with user/collection parameters - processor.vecstore.insert.assert_called_once_with( - [0.1, 0.2, 0.3], "Valid document content", 'test_user', 'test_collection' - ) + + # Verify valid chunks were inserted, empty string chunk was skipped + expected_calls = [ + ([0.1, 0.2, 0.3], "Valid document content", 'test_user', 'test_collection'), + ([0.7, 0.8, 0.9], "Another valid chunk", 'test_user', 'test_collection'), + ] + + assert processor.vecstore.insert.call_count == 2 + for i, (expected_vec, expected_chunk_id, expected_user, expected_collection) in enumerate(expected_calls): + actual_call = processor.vecstore.insert.call_args_list[i] + assert actual_call[0][0] == expected_vec + assert actual_call[0][1] == expected_chunk_id + assert actual_call[0][2] == expected_user + assert actual_call[0][3] == expected_collection @pytest.mark.asyncio async def test_store_document_embeddings_empty_chunks_list(self, processor):