From f35636970849f0c112c14a0f6ccd44f6823d07e2 Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Thu, 15 Jan 2026 12:00:41 +0100 Subject: [PATCH] refactor: Update unit tests --- .../tests/unit/modules/retrieval/chunks_retriever_test.py | 8 ++++++-- .../unit/modules/retrieval/summaries_retriever_test.py | 8 ++++++-- .../unit/modules/retrieval/triplet_retriever_test.py | 4 +++- 3 files changed, 15 insertions(+), 5 deletions(-) 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