diff --git a/cognee/tests/test_kuzu.py b/cognee/tests/test_kuzu.py
index d60c6085e..16c7b9cf6 100644
--- a/cognee/tests/test_kuzu.py
+++ b/cognee/tests/test_kuzu.py
@@ -94,12 +94,12 @@ async def main():
 
         await cognee.cognify([dataset_name])
 
-        context_nonempty = await GraphCompletionRetriever(
+        context_nonempty, _ = await GraphCompletionRetriever(
             node_type=NodeSet,
             node_name=["first"],
         ).get_context("What is in the context?")
 
-        context_empty = await GraphCompletionRetriever(
+        context_empty, _ = await GraphCompletionRetriever(
             node_type=NodeSet,
             node_name=["nonexistent"],
         ).get_context("What is in the context?")
diff --git a/cognee/tests/test_neo4j.py b/cognee/tests/test_neo4j.py
index dcbb38963..d5ccbc19e 100644
--- a/cognee/tests/test_neo4j.py
+++ b/cognee/tests/test_neo4j.py
@@ -98,12 +98,12 @@ async def main():
 
     await cognee.cognify([dataset_name])
 
-    context_nonempty = await GraphCompletionRetriever(
+    context_nonempty, _ = await GraphCompletionRetriever(
         node_type=NodeSet,
         node_name=["first"],
     ).get_context("What is in the context?")
 
-    context_empty = await GraphCompletionRetriever(
+    context_empty, _ = await GraphCompletionRetriever(
         node_type=NodeSet,
         node_name=["nonexistent"],
     ).get_context("What is in the context?")