From eecf1fb11dab8292c98dd1aa2fea47ca90aaace7 Mon Sep 17 00:00:00 2001 From: hajdul88 <52442977+hajdul88@users.noreply.github.com> Date: Thu, 11 Dec 2025 13:23:48 +0100 Subject: [PATCH] adds unit test for graph summary completion retriever --- ...test_graph_summary_completion_retriever.py | 157 ++++++++++++++++++ 1 file changed, 157 insertions(+) create mode 100644 cognee/tests/unit/modules/retrieval/test_graph_summary_completion_retriever.py diff --git a/cognee/tests/unit/modules/retrieval/test_graph_summary_completion_retriever.py b/cognee/tests/unit/modules/retrieval/test_graph_summary_completion_retriever.py new file mode 100644 index 000000000..2af10da5e --- /dev/null +++ b/cognee/tests/unit/modules/retrieval/test_graph_summary_completion_retriever.py @@ -0,0 +1,157 @@ +import pytest +from unittest.mock import AsyncMock, patch, MagicMock + +from cognee.modules.retrieval.graph_summary_completion_retriever import ( + GraphSummaryCompletionRetriever, +) +from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge + + +@pytest.fixture +def mock_edge(): + """Create a mock edge.""" + edge = MagicMock(spec=Edge) + return edge + + +class TestGraphSummaryCompletionRetriever: + @pytest.mark.asyncio + async def test_init_defaults(self): + """Test GraphSummaryCompletionRetriever initialization with defaults.""" + retriever = GraphSummaryCompletionRetriever() + + assert retriever.summarize_prompt_path == "summarize_search_results.txt" + assert retriever.user_prompt_path == "graph_context_for_question.txt" + assert retriever.system_prompt_path == "answer_simple_question.txt" + assert retriever.top_k == 5 + assert retriever.save_interaction is False + + @pytest.mark.asyncio + async def test_init_custom_params(self): + """Test GraphSummaryCompletionRetriever initialization with custom parameters.""" + retriever = GraphSummaryCompletionRetriever( + user_prompt_path="custom_user.txt", + system_prompt_path="custom_system.txt", + summarize_prompt_path="custom_summarize.txt", + system_prompt="Custom system prompt", + top_k=10, + save_interaction=True, + wide_search_top_k=200, + triplet_distance_penalty=2.5, + ) + + assert retriever.summarize_prompt_path == "custom_summarize.txt" + assert retriever.user_prompt_path == "custom_user.txt" + assert retriever.system_prompt_path == "custom_system.txt" + assert retriever.top_k == 10 + assert retriever.save_interaction is True + + @pytest.mark.asyncio + async def test_resolve_edges_to_text_calls_super_and_summarizes(self, mock_edge): + """Test resolve_edges_to_text calls super method and then summarizes.""" + retriever = GraphSummaryCompletionRetriever( + summarize_prompt_path="custom_summarize.txt", + system_prompt="Custom system prompt", + ) + + with ( + patch( + "cognee.modules.retrieval.graph_completion_retriever.GraphCompletionRetriever.resolve_edges_to_text", + new_callable=AsyncMock, + return_value="Resolved edges text", + ) as mock_super_resolve, + patch( + "cognee.modules.retrieval.graph_summary_completion_retriever.summarize_text", + new_callable=AsyncMock, + return_value="Summarized text", + ) as mock_summarize, + ): + result = await retriever.resolve_edges_to_text([mock_edge]) + + assert result == "Summarized text" + mock_super_resolve.assert_awaited_once_with([mock_edge]) + mock_summarize.assert_awaited_once_with( + "Resolved edges text", + "custom_summarize.txt", + "Custom system prompt", + ) + + @pytest.mark.asyncio + async def test_resolve_edges_to_text_with_default_system_prompt(self, mock_edge): + """Test resolve_edges_to_text uses None for system_prompt when not provided.""" + retriever = GraphSummaryCompletionRetriever() + + with ( + patch( + "cognee.modules.retrieval.graph_completion_retriever.GraphCompletionRetriever.resolve_edges_to_text", + new_callable=AsyncMock, + return_value="Resolved edges text", + ), + patch( + "cognee.modules.retrieval.graph_summary_completion_retriever.summarize_text", + new_callable=AsyncMock, + return_value="Summarized text", + ) as mock_summarize, + ): + await retriever.resolve_edges_to_text([mock_edge]) + + mock_summarize.assert_awaited_once_with( + "Resolved edges text", + "summarize_search_results.txt", + None, + ) + + @pytest.mark.asyncio + async def test_resolve_edges_to_text_with_empty_edges(self): + """Test resolve_edges_to_text handles empty edges list.""" + retriever = GraphSummaryCompletionRetriever() + + with ( + patch( + "cognee.modules.retrieval.graph_completion_retriever.GraphCompletionRetriever.resolve_edges_to_text", + new_callable=AsyncMock, + return_value="", + ), + patch( + "cognee.modules.retrieval.graph_summary_completion_retriever.summarize_text", + new_callable=AsyncMock, + return_value="Empty summary", + ) as mock_summarize, + ): + result = await retriever.resolve_edges_to_text([]) + + assert result == "Empty summary" + mock_summarize.assert_awaited_once_with( + "", + "summarize_search_results.txt", + None, + ) + + @pytest.mark.asyncio + async def test_resolve_edges_to_text_with_multiple_edges(self, mock_edge): + """Test resolve_edges_to_text handles multiple edges.""" + retriever = GraphSummaryCompletionRetriever() + + mock_edge2 = MagicMock(spec=Edge) + mock_edge3 = MagicMock(spec=Edge) + + with ( + patch( + "cognee.modules.retrieval.graph_completion_retriever.GraphCompletionRetriever.resolve_edges_to_text", + new_callable=AsyncMock, + return_value="Multiple edges resolved text", + ), + patch( + "cognee.modules.retrieval.graph_summary_completion_retriever.summarize_text", + new_callable=AsyncMock, + return_value="Multiple edges summarized", + ) as mock_summarize, + ): + result = await retriever.resolve_edges_to_text([mock_edge, mock_edge2, mock_edge3]) + + assert result == "Multiple edges summarized" + mock_summarize.assert_awaited_once_with( + "Multiple edges resolved text", + "summarize_search_results.txt", + None, + )