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

refactor: Update unit tests

Browse source
This commit is contained in:
Igor Ilic 2026-01-15 12:00:41 +01:00
parent e65a6a1679
commit f356369708
3 changed files with 15 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
cognee/tests/unit/modules/retrieval/chunks_retriever_test.py
View file

@ -35,7 +35,9 @@ async def test_get_context_success(mock_vector_engine):
assert len(context) == 2
assert context[0]["text"] == "Steve Rodger"
assert context[1]["text"] == "Mike Broski"
mock_vector_engine.search.assert_awaited_once_with("DocumentChunk_text", "test query", limit=5)
mock_vector_engine.search.assert_awaited_once_with(
"DocumentChunk_text", "test query", limit=5, include_payload=True
)
@pytest.mark.asyncio
@ -87,7 +89,9 @@ async def test_get_context_top_k_limit(mock_vector_engine):
context = await retriever.get_context("test query")
assert len(context) == 3
mock_vector_engine.search.assert_awaited_once_with("DocumentChunk_text", "test query", limit=3)
mock_vector_engine.search.assert_awaited_once_with(
"DocumentChunk_text", "test query", limit=3, include_payload=True
)
@pytest.mark.asyncio

8
cognee/tests/unit/modules/retrieval/summaries_retriever_test.py
View file

@ -35,7 +35,9 @@ async def test_get_context_success(mock_vector_engine):
assert len(context) == 2
assert context[0]["text"] == "S.R."
assert context[1]["text"] == "M.B."
mock_vector_engine.search.assert_awaited_once_with("TextSummary_text", "test query", limit=5)
mock_vector_engine.search.assert_awaited_once_with(
"TextSummary_text", "test query", limit=5, include_payload=True
)
@pytest.mark.asyncio
@ -87,7 +89,9 @@ async def test_get_context_top_k_limit(mock_vector_engine):
context = await retriever.get_context("test query")
assert len(context) == 3
mock_vector_engine.search.assert_awaited_once_with("TextSummary_text", "test query", limit=3)
mock_vector_engine.search.assert_awaited_once_with(
"TextSummary_text", "test query", limit=3, include_payload=True
)
@pytest.mark.asyncio

4
cognee/tests/unit/modules/retrieval/triplet_retriever_test.py
View file

@ -34,7 +34,9 @@ async def test_get_context_success(mock_vector_engine):
context = await retriever.get_context("test query")
assert context == "Alice knows Bob\nBob works at Tech Corp"
mock_vector_engine.search.assert_awaited_once_with("Triplet_text", "test query", limit=5)
mock_vector_engine.search.assert_awaited_once_with(
"Triplet_text", "test query", limit=5, include_payload=True
)
@pytest.mark.asyncio