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

refactor: Resolve unit tests failing for search

Browse source
This commit is contained in:
Igor Ilic 2025-08-29 10:30:49 +02:00
parent 966e676d61
commit 5bfae7a36b
2 changed files with 10 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
cognee/modules/search/methods/search.py
View file

@ -179,7 +179,7 @@ async def specific_search(
send_telemetry("cognee.search EXECUTION STARTED", user.id) send_telemetry("cognee.search EXECUTION STARTED", user.id)
results = await search_task(query=query_text) results = await search_task(query_text)
send_telemetry("cognee.search EXECUTION COMPLETED", user.id) send_telemetry("cognee.search EXECUTION COMPLETED", user.id)

13
cognee/tests/unit/modules/search/search_methods_test.py
View file

@ -58,15 +58,17 @@ async def test_search(
# Verify # Verify
mock_log_query.assert_called_once_with(query_text, query_type.value, mock_user.id) mock_log_query.assert_called_once_with(query_text, query_type.value, mock_user.id)
mock_specific_search.assert_called_once_with( mock_specific_search.assert_called_once_with(
query_type, query_type=query_type,
query_text, query_text=query_text,
mock_user, user=mock_user,
system_prompt_path="answer_simple_question.txt", system_prompt_path="answer_simple_question.txt",
system_prompt=None,
top_k=10, top_k=10,
node_type=None, node_type=None,
node_name=None, node_name=None,
save_interaction=False, save_interaction=False,
last_k=None, last_k=None,
only_context=False,
) )
# Verify result logging # Verify result logging
@ -201,7 +203,10 @@ async def test_specific_search_feeling_lucky(
if retriever_name == "CompletionRetriever": if retriever_name == "CompletionRetriever":
mock_retriever_class.assert_called_once_with( mock_retriever_class.assert_called_once_with(
system_prompt_path="answer_simple_question.txt", top_k=top_k system_prompt_path="answer_simple_question.txt",
top_k=top_k,
system_prompt=None,
only_context=None,
) )
else: else:
mock_retriever_class.assert_called_once_with(top_k=top_k) mock_retriever_class.assert_called_once_with(top_k=top_k)