From 8e0d1124399fb6f6592c279c564700b4355732fd Mon Sep 17 00:00:00 2001 From: Andrej Milicevic Date: Mon, 19 Jan 2026 10:56:50 +0100 Subject: [PATCH] test: tiny test fix --- .../graph_completion_retriever_context_extension_test.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/cognee/tests/unit/modules/retrieval/graph_completion_retriever_context_extension_test.py b/cognee/tests/unit/modules/retrieval/graph_completion_retriever_context_extension_test.py index 8d6310214..7ce0d3a57 100644 --- a/cognee/tests/unit/modules/retrieval/graph_completion_retriever_context_extension_test.py +++ b/cognee/tests/unit/modules/retrieval/graph_completion_retriever_context_extension_test.py @@ -611,7 +611,12 @@ async def test_get_completion_batch_queries_context_extension_stops_early(mock_e retriever = GraphCompletionContextExtensionRetriever() with ( - patch.object(retriever, "get_context", new_callable=AsyncMock, return_value=[[mock_edge]]), + patch.object( + retriever, + "get_context", + new_callable=AsyncMock, + return_value=[[mock_edge], [mock_edge]], + ), patch( "cognee.modules.retrieval.graph_completion_retriever.resolve_edges_to_text", return_value="Resolved context",