test: tiny test fix

This commit is contained in:
Andrej Milicevic 2026-01-19 10:56:50 +01:00
parent d258b1d7af
commit 8e0d112439

View file

@ -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",