This commit is contained in:
hajdul88 2025-12-10 18:18:24 +01:00
parent 0329777290
commit ff1add1af0

View file

@ -92,15 +92,17 @@ async def test_get_completion_without_session(mock_vector_engine):
retriever = TripletRetriever() retriever = TripletRetriever()
with patch( with (
"cognee.modules.retrieval.triplet_retriever.get_vector_engine", patch(
return_value=mock_vector_engine, "cognee.modules.retrieval.triplet_retriever.get_vector_engine",
), patch( return_value=mock_vector_engine,
"cognee.modules.retrieval.triplet_retriever.generate_completion", ),
return_value="Generated answer", patch(
), patch( "cognee.modules.retrieval.triplet_retriever.generate_completion",
"cognee.modules.retrieval.triplet_retriever.CacheConfig" return_value="Generated answer",
) as mock_cache_config: ),
patch("cognee.modules.retrieval.triplet_retriever.CacheConfig") as mock_cache_config,
):
mock_config = MagicMock() mock_config = MagicMock()
mock_config.caching = False mock_config.caching = False
mock_cache_config.return_value = mock_config mock_cache_config.return_value = mock_config
@ -117,12 +119,13 @@ async def test_get_completion_with_provided_context(mock_vector_engine):
"""Test get_completion with provided context.""" """Test get_completion with provided context."""
retriever = TripletRetriever() retriever = TripletRetriever()
with patch( with (
"cognee.modules.retrieval.triplet_retriever.generate_completion", patch(
return_value="Generated answer", "cognee.modules.retrieval.triplet_retriever.generate_completion",
), patch( return_value="Generated answer",
"cognee.modules.retrieval.triplet_retriever.CacheConfig" ),
) as mock_cache_config: patch("cognee.modules.retrieval.triplet_retriever.CacheConfig") as mock_cache_config,
):
mock_config = MagicMock() mock_config = MagicMock()
mock_config.caching = False mock_config.caching = False
mock_cache_config.return_value = mock_config mock_cache_config.return_value = mock_config