diff --git a/tests/unit/test_retrieval/test_graph_rag.py b/tests/unit/test_retrieval/test_graph_rag.py new file mode 100644 index 00000000..3afe1a4f --- /dev/null +++ b/tests/unit/test_retrieval/test_graph_rag.py @@ -0,0 +1,61 @@ +""" +Tests for GraphRAG retrieval implementation +""" + +import pytest +from unittest.mock import MagicMock + +from trustgraph.retrieval.graph_rag.graph_rag import GraphRag + + +class TestGraphRag: + """Test cases for GraphRag class""" + + def test_graph_rag_initialization_with_defaults(self): + """Test GraphRag initialization with default verbose setting""" + # Create mock clients + mock_prompt_client = MagicMock() + mock_embeddings_client = MagicMock() + mock_graph_embeddings_client = MagicMock() + mock_triples_client = MagicMock() + + # Initialize GraphRag + graph_rag = GraphRag( + prompt_client=mock_prompt_client, + embeddings_client=mock_embeddings_client, + graph_embeddings_client=mock_graph_embeddings_client, + triples_client=mock_triples_client + ) + + # Verify initialization + assert graph_rag.prompt_client == mock_prompt_client + assert graph_rag.embeddings_client == mock_embeddings_client + assert graph_rag.graph_embeddings_client == mock_graph_embeddings_client + assert graph_rag.triples_client == mock_triples_client + assert graph_rag.verbose is False # Default value + assert graph_rag.label_cache == {} # Empty cache initially + + def test_graph_rag_initialization_with_verbose(self): + """Test GraphRag initialization with verbose enabled""" + # Create mock clients + mock_prompt_client = MagicMock() + mock_embeddings_client = MagicMock() + mock_graph_embeddings_client = MagicMock() + mock_triples_client = MagicMock() + + # Initialize GraphRag with verbose=True + graph_rag = GraphRag( + prompt_client=mock_prompt_client, + embeddings_client=mock_embeddings_client, + graph_embeddings_client=mock_graph_embeddings_client, + triples_client=mock_triples_client, + verbose=True + ) + + # Verify initialization + assert graph_rag.prompt_client == mock_prompt_client + assert graph_rag.embeddings_client == mock_embeddings_client + assert graph_rag.graph_embeddings_client == mock_graph_embeddings_client + assert graph_rag.triples_client == mock_triples_client + assert graph_rag.verbose is True + assert graph_rag.label_cache == {} # Empty cache initially \ No newline at end of file