From fe82b844c6b8bee232bab32676febbbe27b444db Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Mon, 9 Mar 2026 10:09:50 +0000 Subject: [PATCH] Fixing tests --- tests/integration/test_document_rag_integration.py | 6 +++--- tests/integration/test_graph_rag_integration.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/integration/test_document_rag_integration.py b/tests/integration/test_document_rag_integration.py index 80520429..997f5f61 100644 --- a/tests/integration/test_document_rag_integration.py +++ b/tests/integration/test_document_rag_integration.py @@ -312,9 +312,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 chunk_id set (100 chunks) - large_chunk_ids = [f"doc/c{i}" for i in range(100)] - mock_doc_embeddings_client.query.return_value = large_chunk_ids + # Arrange - Mock large chunk match set (100 chunks) + large_chunk_matches = [ChunkMatch(chunk_id=f"doc/c{i}", score=0.9 - i*0.001) for i in range(100)] + mock_doc_embeddings_client.query.return_value = large_chunk_matches # Act import time diff --git a/tests/integration/test_graph_rag_integration.py b/tests/integration/test_graph_rag_integration.py index d3b6c2ba..260b93d5 100644 --- a/tests/integration/test_graph_rag_integration.py +++ b/tests/integration/test_graph_rag_integration.py @@ -131,7 +131,7 @@ class TestGraphRagIntegration: # 2. Should query graph embeddings to find relevant entities mock_graph_embeddings_client.query.assert_called_once() call_args = mock_graph_embeddings_client.query.call_args - assert call_args.kwargs['vectors'] == [[0.1, 0.2, 0.3, 0.4, 0.5]] + assert call_args.kwargs['vector'] == [0.1, 0.2, 0.3, 0.4, 0.5] assert call_args.kwargs['limit'] == entity_limit assert call_args.kwargs['user'] == user assert call_args.kwargs['collection'] == collection