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

ruff

Browse source
This commit is contained in:
hajdul88 2025-12-11 12:27:27 +01:00
parent c5c60ccad0
commit 0d0fe3aa30
1 changed files with 35 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

44
cognee/tests/unit/modules/retrieval/temporal_retriever_test.py
View file

@ -616,9 +616,15 @@ async def test_extract_time_from_query_relative_path():
with (
patch("cognee.modules.retrieval.temporal_retriever.os.path.isabs", return_value=False),
patch("cognee.modules.retrieval.temporal_retriever.datetime") as mock_datetime,
patch("cognee.modules.retrieval.temporal_retriever.render_prompt", return_value="System prompt"),
patch(
"cognee.modules.retrieval.temporal_retriever.render_prompt",
return_value="System prompt",
),
patch.object(
LLMGateway, "acreate_structured_output", new_callable=AsyncMock, return_value=mock_interval
LLMGateway,
"acreate_structured_output",
new_callable=AsyncMock,
return_value=mock_interval,
),
):
mock_datetime.now.return_value.strftime.return_value = "11-12-2024"
@ -632,7 +638,9 @@ async def test_extract_time_from_query_relative_path():
@pytest.mark.asyncio
async def test_extract_time_from_query_absolute_path():
"""Test extract_time_from_query with absolute prompt path."""
retriever = TemporalRetriever(time_extraction_prompt_path="/absolute/path/to/extract_query_time.txt")
retriever = TemporalRetriever(
time_extraction_prompt_path="/absolute/path/to/extract_query_time.txt"
)
mock_timestamp_from = Timestamp(year=2024, month=1, day=1)
mock_timestamp_to = Timestamp(year=2024, month=12, day=31)
@ -640,12 +648,24 @@ async def test_extract_time_from_query_absolute_path():
with (
patch("cognee.modules.retrieval.temporal_retriever.os.path.isabs", return_value=True),
patch("cognee.modules.retrieval.temporal_retriever.os.path.dirname", return_value="/absolute/path/to"),
patch("cognee.modules.retrieval.temporal_retriever.os.path.basename", return_value="extract_query_time.txt"),
patch(
"cognee.modules.retrieval.temporal_retriever.os.path.dirname",
return_value="/absolute/path/to",
),
patch(
"cognee.modules.retrieval.temporal_retriever.os.path.basename",
return_value="extract_query_time.txt",
),
patch("cognee.modules.retrieval.temporal_retriever.datetime") as mock_datetime,
patch("cognee.modules.retrieval.temporal_retriever.render_prompt", return_value="System prompt"),
patch(
"cognee.modules.retrieval.temporal_retriever.render_prompt",
return_value="System prompt",
),
patch.object(
LLMGateway, "acreate_structured_output", new_callable=AsyncMock, return_value=mock_interval
LLMGateway,
"acreate_structured_output",
new_callable=AsyncMock,
return_value=mock_interval,
),
):
mock_datetime.now.return_value.strftime.return_value = "11-12-2024"
@ -666,9 +686,15 @@ async def test_extract_time_from_query_with_none_values():
with (
patch("cognee.modules.retrieval.temporal_retriever.os.path.isabs", return_value=False),
patch("cognee.modules.retrieval.temporal_retriever.datetime") as mock_datetime,
patch("cognee.modules.retrieval.temporal_retriever.render_prompt", return_value="System prompt"),
patch(
"cognee.modules.retrieval.temporal_retriever.render_prompt",
return_value="System prompt",
),
patch.object(
LLMGateway, "acreate_structured_output", new_callable=AsyncMock, return_value=mock_interval
LLMGateway,
"acreate_structured_output",
new_callable=AsyncMock,
return_value=mock_interval,
),
):
mock_datetime.now.return_value.strftime.return_value = "11-12-2024"