From 78fb4158927e9c5cd6b7a0fa29bb7dafd23a1a93 Mon Sep 17 00:00:00 2001
From: hajdul88 <52442977+hajdul88@users.noreply.github.com>
Date: Mon, 18 Aug 2025 13:40:33 +0200
Subject: [PATCH] chore: changes context  return value in tests

---
 .../graph_completion_retriever_context_extension_test.py    | 6 +++---
 .../retrieval/graph_completion_retriever_cot_test.py        | 6 +++---
 .../modules/retrieval/graph_completion_retriever_test.py    | 6 +++---
 3 files changed, 9 insertions(+), 9 deletions(-)

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 3e7f6626f..26ae2f883 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
@@ -51,7 +51,7 @@ class TestGraphCompletionWithContextExtensionRetriever:
 
         retriever = GraphCompletionContextExtensionRetriever()
 
-        context = await retriever.get_context("Who works at Canva?")
+        context, _ = await retriever.get_context("Who works at Canva?")
 
         assert "Mike Broski --[works_for]--> Canva" in context, "Failed to get Mike Broski"
         assert "Christina Mayer --[works_for]--> Canva" in context, "Failed to get Christina Mayer"
@@ -129,7 +129,7 @@ class TestGraphCompletionWithContextExtensionRetriever:
 
         retriever = GraphCompletionContextExtensionRetriever(top_k=20)
 
-        context = await retriever.get_context("Who works at Figma?")
+        context, _ = await retriever.get_context("Who works at Figma?")
 
         print(context)
 
@@ -167,7 +167,7 @@ class TestGraphCompletionWithContextExtensionRetriever:
 
         await setup()
 
-        context = await retriever.get_context("Who works at Figma?")
+        context, _ = await retriever.get_context("Who works at Figma?")
         assert context == "", "Context should be empty on an empty graph"
 
         answer = await retriever.get_completion("Who works at Figma?")
diff --git a/cognee/tests/unit/modules/retrieval/graph_completion_retriever_cot_test.py b/cognee/tests/unit/modules/retrieval/graph_completion_retriever_cot_test.py
index ff92dfd8f..be25299aa 100644
--- a/cognee/tests/unit/modules/retrieval/graph_completion_retriever_cot_test.py
+++ b/cognee/tests/unit/modules/retrieval/graph_completion_retriever_cot_test.py
@@ -47,7 +47,7 @@ class TestGraphCompletionCoTRetriever:
 
         retriever = GraphCompletionCotRetriever()
 
-        context = await retriever.get_context("Who works at Canva?")
+        context, _ = await retriever.get_context("Who works at Canva?")
 
         assert "Mike Broski --[works_for]--> Canva" in context, "Failed to get Mike Broski"
         assert "Christina Mayer --[works_for]--> Canva" in context, "Failed to get Christina Mayer"
@@ -124,7 +124,7 @@ class TestGraphCompletionCoTRetriever:
 
         retriever = GraphCompletionCotRetriever(top_k=20)
 
-        context = await retriever.get_context("Who works at Figma?")
+        context, _ = await retriever.get_context("Who works at Figma?")
 
         print(context)
 
@@ -162,7 +162,7 @@ class TestGraphCompletionCoTRetriever:
 
         await setup()
 
-        context = await retriever.get_context("Who works at Figma?")
+        context, _ = await retriever.get_context("Who works at Figma?")
         assert context == "", "Context should be empty on an empty graph"
 
         answer = await retriever.get_completion("Who works at Figma?")
diff --git a/cognee/tests/unit/modules/retrieval/graph_completion_retriever_test.py b/cognee/tests/unit/modules/retrieval/graph_completion_retriever_test.py
index 976b69e69..50784f94a 100644
--- a/cognee/tests/unit/modules/retrieval/graph_completion_retriever_test.py
+++ b/cognee/tests/unit/modules/retrieval/graph_completion_retriever_test.py
@@ -67,7 +67,7 @@ class TestGraphCompletionRetriever:
 
         retriever = GraphCompletionRetriever()
 
-        context = await retriever.get_context("Who works at Canva?")
+        context, _ = await retriever.get_context("Who works at Canva?")
 
         # Ensure the top-level sections are present
         assert "Nodes:" in context, "Missing 'Nodes:' section in context"
@@ -191,7 +191,7 @@ class TestGraphCompletionRetriever:
 
         retriever = GraphCompletionRetriever(top_k=20)
 
-        context = await retriever.get_context("Who works at Figma?")
+        context, _ = await retriever.get_context("Who works at Figma?")
 
         print(context)
 
@@ -222,5 +222,5 @@ class TestGraphCompletionRetriever:
 
         await setup()
 
-        context = await retriever.get_context("Who works at Figma?")
+        context, _ = await retriever.get_context("Who works at Figma?")
         assert context == "", "Context should be empty on an empty graph"