diff --git a/cognee/tests/unit/modules/retrieval/chunks_retriever_test.py b/cognee/tests/unit/modules/retrieval/chunks_retriever_test.py
index 98bfd48fe..feb254155 100644
--- a/cognee/tests/unit/modules/retrieval/chunks_retriever_test.py
+++ b/cognee/tests/unit/modules/retrieval/chunks_retriever_test.py
@@ -35,7 +35,9 @@ async def test_get_context_success(mock_vector_engine):
     assert len(context) == 2
     assert context[0]["text"] == "Steve Rodger"
     assert context[1]["text"] == "Mike Broski"
-    mock_vector_engine.search.assert_awaited_once_with("DocumentChunk_text", "test query", limit=5)
+    mock_vector_engine.search.assert_awaited_once_with(
+        "DocumentChunk_text", "test query", limit=5, include_payload=True
+    )
 
 
 @pytest.mark.asyncio
@@ -87,7 +89,9 @@ async def test_get_context_top_k_limit(mock_vector_engine):
         context = await retriever.get_context("test query")
 
     assert len(context) == 3
-    mock_vector_engine.search.assert_awaited_once_with("DocumentChunk_text", "test query", limit=3)
+    mock_vector_engine.search.assert_awaited_once_with(
+        "DocumentChunk_text", "test query", limit=3, include_payload=True
+    )
 
 
 @pytest.mark.asyncio
diff --git a/cognee/tests/unit/modules/retrieval/summaries_retriever_test.py b/cognee/tests/unit/modules/retrieval/summaries_retriever_test.py
index e552ac74a..7bec8afdf 100644
--- a/cognee/tests/unit/modules/retrieval/summaries_retriever_test.py
+++ b/cognee/tests/unit/modules/retrieval/summaries_retriever_test.py
@@ -35,7 +35,9 @@ async def test_get_context_success(mock_vector_engine):
     assert len(context) == 2
     assert context[0]["text"] == "S.R."
     assert context[1]["text"] == "M.B."
-    mock_vector_engine.search.assert_awaited_once_with("TextSummary_text", "test query", limit=5)
+    mock_vector_engine.search.assert_awaited_once_with(
+        "TextSummary_text", "test query", limit=5, include_payload=True
+    )
 
 
 @pytest.mark.asyncio
@@ -87,7 +89,9 @@ async def test_get_context_top_k_limit(mock_vector_engine):
         context = await retriever.get_context("test query")
 
     assert len(context) == 3
-    mock_vector_engine.search.assert_awaited_once_with("TextSummary_text", "test query", limit=3)
+    mock_vector_engine.search.assert_awaited_once_with(
+        "TextSummary_text", "test query", limit=3, include_payload=True
+    )
 
 
 @pytest.mark.asyncio
diff --git a/cognee/tests/unit/modules/retrieval/triplet_retriever_test.py b/cognee/tests/unit/modules/retrieval/triplet_retriever_test.py
index 83612c7aa..e914b0aa4 100644
--- a/cognee/tests/unit/modules/retrieval/triplet_retriever_test.py
+++ b/cognee/tests/unit/modules/retrieval/triplet_retriever_test.py
@@ -34,7 +34,9 @@ async def test_get_context_success(mock_vector_engine):
         context = await retriever.get_context("test query")
 
     assert context == "Alice knows Bob\nBob works at Tech Corp"
-    mock_vector_engine.search.assert_awaited_once_with("Triplet_text", "test query", limit=5)
+    mock_vector_engine.search.assert_awaited_once_with(
+        "Triplet_text", "test query", limit=5, include_payload=True
+    )
 
 
 @pytest.mark.asyncio