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