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

ruff ruff

Browse source
This commit is contained in:
hajdul88 2025-12-10 18:07:22 +01:00
parent 508086c513
commit ac58058eaa
1 changed files with 38 additions and 24 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -263,14 +263,18 @@ async def test_get_context_with_time_range(mock_graph_engine, mock_vector_engine
mock_result2 = SimpleNamespace(payload={"id": "e1"}, score=0.10) mock_result2 = SimpleNamespace(payload={"id": "e1"}, score=0.10)
mock_vector_engine.search.return_value = [mock_result1, mock_result2] mock_vector_engine.search.return_value = [mock_result1, mock_result2]
with patch.object( with (
retriever, "extract_time_from_query", return_value=("2024-01-01", "2024-12-31") patch.object(
), patch( retriever, "extract_time_from_query", return_value=("2024-01-01", "2024-12-31")
"cognee.modules.retrieval.temporal_retriever.get_graph_engine", ),
return_value=mock_graph_engine, patch(
), patch( "cognee.modules.retrieval.temporal_retriever.get_graph_engine",
"cognee.modules.retrieval.temporal_retriever.get_vector_engine", return_value=mock_graph_engine,
return_value=mock_vector_engine, ),
patch(
"cognee.modules.retrieval.temporal_retriever.get_vector_engine",
return_value=mock_vector_engine,
),
): ):
context = await retriever.get_context("What happened in 2024?") context = await retriever.get_context("What happened in 2024?")
@ -284,14 +288,19 @@ async def test_get_context_fallback_to_triplets_no_time(mock_graph_engine):
"""Test get_context falls back to triplets when no time is extracted.""" """Test get_context falls back to triplets when no time is extracted."""
retriever = TemporalRetriever() retriever = TemporalRetriever()
with patch( with (
"cognee.modules.retrieval.temporal_retriever.get_graph_engine", patch(
return_value=mock_graph_engine, "cognee.modules.retrieval.temporal_retriever.get_graph_engine",
), patch.object( return_value=mock_graph_engine,
retriever, "get_triplets", return_value=[{"s": "a", "p": "b", "o": "c"}] ),
) as mock_get_triplets, patch.object( patch.object(
retriever, "resolve_edges_to_text", return_value="triplet text" retriever, "get_triplets", return_value=[{"s": "a", "p": "b", "o": "c"}]
) as mock_resolve: ) as mock_get_triplets,
patch.object(
retriever, "resolve_edges_to_text", return_value="triplet text"
) as mock_resolve,
):
async def mock_extract_time(query): async def mock_extract_time(query):
return None, None return None, None
@ -311,14 +320,19 @@ async def test_get_context_no_events_found(mock_graph_engine):
mock_graph_engine.collect_time_ids.return_value = [] mock_graph_engine.collect_time_ids.return_value = []
with patch( with (
"cognee.modules.retrieval.temporal_retriever.get_graph_engine", patch(
return_value=mock_graph_engine, "cognee.modules.retrieval.temporal_retriever.get_graph_engine",
), patch.object( return_value=mock_graph_engine,
retriever, "get_triplets", return_value=[{"s": "a", "p": "b", "o": "c"}] ),
) as mock_get_triplets, patch.object( patch.object(
retriever, "resolve_edges_to_text", return_value="triplet text" retriever, "get_triplets", return_value=[{"s": "a", "p": "b", "o": "c"}]
) as mock_resolve: ) as mock_get_triplets,
patch.object(
retriever, "resolve_edges_to_text", return_value="triplet text"
) as mock_resolve,
):
async def mock_extract_time(query): async def mock_extract_time(query):
return "2024-01-01", "2024-12-31" return "2024-01-01", "2024-12-31"