gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

chore: changes context return value in tests

Browse source
This commit is contained in:
hajdul88 2025-08-18 13:40:33 +02:00
parent 1d63da7923
commit 78fb415892
3 changed files with 9 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
cognee/tests/unit/modules/retrieval/graph_completion_retriever_context_extension_test.py
View file

@ -51,7 +51,7 @@ class TestGraphCompletionWithContextExtensionRetriever:
retriever = GraphCompletionContextExtensionRetriever() 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 "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" assert "Christina Mayer --[works_for]--> Canva" in context, "Failed to get Christina Mayer"
@ -129,7 +129,7 @@ class TestGraphCompletionWithContextExtensionRetriever:
retriever = GraphCompletionContextExtensionRetriever(top_k=20) 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) print(context)
@ -167,7 +167,7 @@ class TestGraphCompletionWithContextExtensionRetriever:
await setup() 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" assert context == "", "Context should be empty on an empty graph"
answer = await retriever.get_completion("Who works at Figma?") answer = await retriever.get_completion("Who works at Figma?")

6
cognee/tests/unit/modules/retrieval/graph_completion_retriever_cot_test.py
View file

@ -47,7 +47,7 @@ class TestGraphCompletionCoTRetriever:
retriever = GraphCompletionCotRetriever() 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 "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" assert "Christina Mayer --[works_for]--> Canva" in context, "Failed to get Christina Mayer"
@ -124,7 +124,7 @@ class TestGraphCompletionCoTRetriever:
retriever = GraphCompletionCotRetriever(top_k=20) 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) print(context)
@ -162,7 +162,7 @@ class TestGraphCompletionCoTRetriever:
await setup() 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" assert context == "", "Context should be empty on an empty graph"
answer = await retriever.get_completion("Who works at Figma?") answer = await retriever.get_completion("Who works at Figma?")

6
cognee/tests/unit/modules/retrieval/graph_completion_retriever_test.py
View file

@ -67,7 +67,7 @@ class TestGraphCompletionRetriever:
retriever = GraphCompletionRetriever() 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 # Ensure the top-level sections are present
assert "Nodes:" in context, "Missing 'Nodes:' section in context" assert "Nodes:" in context, "Missing 'Nodes:' section in context"
@ -191,7 +191,7 @@ class TestGraphCompletionRetriever:
retriever = GraphCompletionRetriever(top_k=20) 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) print(context)
@ -222,5 +222,5 @@ class TestGraphCompletionRetriever:
await setup() 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" assert context == "", "Context should be empty on an empty graph"