From b1d05250746c91e83bb0fcb68a11d95f203eafc7 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Sat, 12 Jul 2025 11:12:19 +0100 Subject: [PATCH] GraphRAG testing --- tests/unit/test_retrieval/test_graph_rag.py | 540 +++++++++++++++++++- 1 file changed, 537 insertions(+), 3 deletions(-) diff --git a/tests/unit/test_retrieval/test_graph_rag.py b/tests/unit/test_retrieval/test_graph_rag.py index 3afe1a4f..788f71a2 100644 --- a/tests/unit/test_retrieval/test_graph_rag.py +++ b/tests/unit/test_retrieval/test_graph_rag.py @@ -3,9 +3,10 @@ Tests for GraphRAG retrieval implementation """ import pytest -from unittest.mock import MagicMock +import unittest.mock +from unittest.mock import MagicMock, AsyncMock -from trustgraph.retrieval.graph_rag.graph_rag import GraphRag +from trustgraph.retrieval.graph_rag.graph_rag import GraphRag, Query class TestGraphRag: @@ -58,4 +59,537 @@ class TestGraphRag: 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 + assert graph_rag.label_cache == {} # Empty cache initially + + +class TestQuery: + """Test cases for Query class""" + + def test_query_initialization_with_defaults(self): + """Test Query initialization with default parameters""" + # Create mock GraphRag + mock_rag = MagicMock() + + # Initialize Query with defaults + query = Query( + rag=mock_rag, + user="test_user", + collection="test_collection", + verbose=False + ) + + # Verify initialization + assert query.rag == mock_rag + assert query.user == "test_user" + assert query.collection == "test_collection" + assert query.verbose is False + assert query.entity_limit == 50 # Default value + assert query.triple_limit == 30 # Default value + assert query.max_subgraph_size == 1000 # Default value + assert query.max_path_length == 2 # Default value + + def test_query_initialization_with_custom_params(self): + """Test Query initialization with custom parameters""" + # Create mock GraphRag + mock_rag = MagicMock() + + # Initialize Query with custom parameters + query = Query( + rag=mock_rag, + user="custom_user", + collection="custom_collection", + verbose=True, + entity_limit=100, + triple_limit=60, + max_subgraph_size=2000, + max_path_length=3 + ) + + # Verify initialization + assert query.rag == mock_rag + assert query.user == "custom_user" + assert query.collection == "custom_collection" + assert query.verbose is True + assert query.entity_limit == 100 + assert query.triple_limit == 60 + assert query.max_subgraph_size == 2000 + assert query.max_path_length == 3 + + @pytest.mark.asyncio + async def test_get_vector_method(self): + """Test Query.get_vector method calls embeddings client correctly""" + # Create mock GraphRag with embeddings client + mock_rag = MagicMock() + mock_embeddings_client = AsyncMock() + mock_rag.embeddings_client = mock_embeddings_client + + # Mock the embed method to return test vectors + expected_vectors = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]] + mock_embeddings_client.embed.return_value = expected_vectors + + # Initialize Query + query = Query( + rag=mock_rag, + user="test_user", + collection="test_collection", + verbose=False + ) + + # Call get_vector + test_query = "What is the capital of France?" + result = await query.get_vector(test_query) + + # Verify embeddings client was called correctly + mock_embeddings_client.embed.assert_called_once_with(test_query) + + # Verify result matches expected vectors + assert result == expected_vectors + + @pytest.mark.asyncio + async def test_get_vector_method_with_verbose(self): + """Test Query.get_vector method with verbose output""" + # Create mock GraphRag with embeddings client + mock_rag = MagicMock() + mock_embeddings_client = AsyncMock() + mock_rag.embeddings_client = mock_embeddings_client + + # Mock the embed method + expected_vectors = [[0.7, 0.8, 0.9]] + mock_embeddings_client.embed.return_value = expected_vectors + + # Initialize Query with verbose=True + query = Query( + rag=mock_rag, + user="test_user", + collection="test_collection", + verbose=True + ) + + # Call get_vector + test_query = "Test query for embeddings" + result = await query.get_vector(test_query) + + # Verify embeddings client was called correctly + mock_embeddings_client.embed.assert_called_once_with(test_query) + + # Verify result matches expected vectors + assert result == expected_vectors + + @pytest.mark.asyncio + async def test_get_entities_method(self): + """Test Query.get_entities method retrieves entities correctly""" + # Create mock GraphRag with clients + mock_rag = MagicMock() + mock_embeddings_client = AsyncMock() + mock_graph_embeddings_client = AsyncMock() + mock_rag.embeddings_client = mock_embeddings_client + mock_rag.graph_embeddings_client = mock_graph_embeddings_client + + # Mock the embedding and entity query responses + test_vectors = [[0.1, 0.2, 0.3]] + mock_embeddings_client.embed.return_value = test_vectors + + # Mock entity objects that have string representation + mock_entity1 = MagicMock() + mock_entity1.__str__ = MagicMock(return_value="entity1") + mock_entity2 = MagicMock() + mock_entity2.__str__ = MagicMock(return_value="entity2") + mock_graph_embeddings_client.query.return_value = [mock_entity1, mock_entity2] + + # Initialize Query + query = Query( + rag=mock_rag, + user="test_user", + collection="test_collection", + verbose=False, + entity_limit=25 + ) + + # Call get_entities + test_query = "Find related entities" + result = await query.get_entities(test_query) + + # Verify embeddings client was called + mock_embeddings_client.embed.assert_called_once_with(test_query) + + # Verify graph embeddings client was called correctly + mock_graph_embeddings_client.query.assert_called_once_with( + vectors=test_vectors, + limit=25, + user="test_user", + collection="test_collection" + ) + + # Verify result is list of entity strings + assert result == ["entity1", "entity2"] + + @pytest.mark.asyncio + async def test_maybe_label_with_cached_label(self): + """Test Query.maybe_label method with cached label""" + # Create mock GraphRag with label cache + mock_rag = MagicMock() + mock_rag.label_cache = {"entity1": "Entity One Label"} + + # Initialize Query + query = Query( + rag=mock_rag, + user="test_user", + collection="test_collection", + verbose=False + ) + + # Call maybe_label with cached entity + result = await query.maybe_label("entity1") + + # Verify cached label is returned + assert result == "Entity One Label" + + @pytest.mark.asyncio + async def test_maybe_label_with_label_lookup(self): + """Test Query.maybe_label method with database label lookup""" + # Create mock GraphRag with triples client + mock_rag = MagicMock() + mock_rag.label_cache = {} # Empty cache + mock_triples_client = AsyncMock() + mock_rag.triples_client = mock_triples_client + + # Mock triple result with label + mock_triple = MagicMock() + mock_triple.o = "Human Readable Label" + mock_triples_client.query.return_value = [mock_triple] + + # Initialize Query + query = Query( + rag=mock_rag, + user="test_user", + collection="test_collection", + verbose=False + ) + + # Call maybe_label + result = await query.maybe_label("http://example.com/entity") + + # Verify triples client was called correctly + mock_triples_client.query.assert_called_once_with( + s="http://example.com/entity", + p="http://www.w3.org/2000/01/rdf-schema#label", + o=None, + limit=1, + user="test_user", + collection="test_collection" + ) + + # Verify result and cache update + assert result == "Human Readable Label" + assert mock_rag.label_cache["http://example.com/entity"] == "Human Readable Label" + + @pytest.mark.asyncio + async def test_maybe_label_with_no_label_found(self): + """Test Query.maybe_label method when no label is found""" + # Create mock GraphRag with triples client + mock_rag = MagicMock() + mock_rag.label_cache = {} # Empty cache + mock_triples_client = AsyncMock() + mock_rag.triples_client = mock_triples_client + + # Mock empty result (no label found) + mock_triples_client.query.return_value = [] + + # Initialize Query + query = Query( + rag=mock_rag, + user="test_user", + collection="test_collection", + verbose=False + ) + + # Call maybe_label + result = await query.maybe_label("unlabeled_entity") + + # Verify triples client was called + mock_triples_client.query.assert_called_once_with( + s="unlabeled_entity", + p="http://www.w3.org/2000/01/rdf-schema#label", + o=None, + limit=1, + user="test_user", + collection="test_collection" + ) + + # Verify result is entity itself and cache is updated + assert result == "unlabeled_entity" + assert mock_rag.label_cache["unlabeled_entity"] == "unlabeled_entity" + + @pytest.mark.asyncio + async def test_follow_edges_basic_functionality(self): + """Test Query.follow_edges method basic triple discovery""" + # Create mock GraphRag with triples client + mock_rag = MagicMock() + mock_triples_client = AsyncMock() + mock_rag.triples_client = mock_triples_client + + # Mock triple results for different query patterns + mock_triple1 = MagicMock() + mock_triple1.s, mock_triple1.p, mock_triple1.o = "entity1", "predicate1", "object1" + + mock_triple2 = MagicMock() + mock_triple2.s, mock_triple2.p, mock_triple2.o = "subject2", "entity1", "object2" + + mock_triple3 = MagicMock() + mock_triple3.s, mock_triple3.p, mock_triple3.o = "subject3", "predicate3", "entity1" + + # Setup query responses for s=ent, p=ent, o=ent patterns + mock_triples_client.query.side_effect = [ + [mock_triple1], # s=ent, p=None, o=None + [mock_triple2], # s=None, p=ent, o=None + [mock_triple3], # s=None, p=None, o=ent + ] + + # Initialize Query + query = Query( + rag=mock_rag, + user="test_user", + collection="test_collection", + verbose=False, + triple_limit=10 + ) + + # Call follow_edges + subgraph = set() + await query.follow_edges("entity1", subgraph, path_length=1) + + # Verify all three query patterns were called + assert mock_triples_client.query.call_count == 3 + + # Verify query calls + mock_triples_client.query.assert_any_call( + s="entity1", p=None, o=None, limit=10, + user="test_user", collection="test_collection" + ) + mock_triples_client.query.assert_any_call( + s=None, p="entity1", o=None, limit=10, + user="test_user", collection="test_collection" + ) + mock_triples_client.query.assert_any_call( + s=None, p=None, o="entity1", limit=10, + user="test_user", collection="test_collection" + ) + + # Verify subgraph contains discovered triples + expected_subgraph = { + ("entity1", "predicate1", "object1"), + ("subject2", "entity1", "object2"), + ("subject3", "predicate3", "entity1") + } + assert subgraph == expected_subgraph + + @pytest.mark.asyncio + async def test_follow_edges_with_path_length_zero(self): + """Test Query.follow_edges method with path_length=0""" + # Create mock GraphRag + mock_rag = MagicMock() + mock_triples_client = AsyncMock() + mock_rag.triples_client = mock_triples_client + + # Initialize Query + query = Query( + rag=mock_rag, + user="test_user", + collection="test_collection", + verbose=False + ) + + # Call follow_edges with path_length=0 + subgraph = set() + await query.follow_edges("entity1", subgraph, path_length=0) + + # Verify no queries were made + mock_triples_client.query.assert_not_called() + + # Verify subgraph remains empty + assert subgraph == set() + + @pytest.mark.asyncio + async def test_follow_edges_with_max_subgraph_size_limit(self): + """Test Query.follow_edges method respects max_subgraph_size""" + # Create mock GraphRag + mock_rag = MagicMock() + mock_triples_client = AsyncMock() + mock_rag.triples_client = mock_triples_client + + # Initialize Query with small max_subgraph_size + query = Query( + rag=mock_rag, + user="test_user", + collection="test_collection", + verbose=False, + max_subgraph_size=2 + ) + + # Pre-populate subgraph to exceed limit + subgraph = {("s1", "p1", "o1"), ("s2", "p2", "o2"), ("s3", "p3", "o3")} + + # Call follow_edges + await query.follow_edges("entity1", subgraph, path_length=1) + + # Verify no queries were made due to size limit + mock_triples_client.query.assert_not_called() + + # Verify subgraph unchanged + assert len(subgraph) == 3 + + @pytest.mark.asyncio + async def test_get_subgraph_method(self): + """Test Query.get_subgraph method orchestrates entity and edge discovery""" + # Create mock Query that patches get_entities and follow_edges + mock_rag = MagicMock() + + query = Query( + rag=mock_rag, + user="test_user", + collection="test_collection", + verbose=False, + max_path_length=1 + ) + + # Mock get_entities to return test entities + query.get_entities = AsyncMock(return_value=["entity1", "entity2"]) + + # Mock follow_edges to add triples to subgraph + async def mock_follow_edges(ent, subgraph, path_length): + subgraph.add((ent, "predicate", "object")) + + query.follow_edges = AsyncMock(side_effect=mock_follow_edges) + + # Call get_subgraph + result = await query.get_subgraph("test query") + + # Verify get_entities was called + query.get_entities.assert_called_once_with("test query") + + # Verify follow_edges was called for each entity + assert query.follow_edges.call_count == 2 + query.follow_edges.assert_any_call("entity1", unittest.mock.ANY, 1) + query.follow_edges.assert_any_call("entity2", unittest.mock.ANY, 1) + + # Verify result is list format + assert isinstance(result, list) + assert len(result) == 2 + + @pytest.mark.asyncio + async def test_get_labelgraph_method(self): + """Test Query.get_labelgraph method converts entities to labels""" + # Create mock Query + mock_rag = MagicMock() + + query = Query( + rag=mock_rag, + user="test_user", + collection="test_collection", + verbose=False, + max_subgraph_size=100 + ) + + # Mock get_subgraph to return test triples + test_subgraph = [ + ("entity1", "predicate1", "object1"), + ("subject2", "http://www.w3.org/2000/01/rdf-schema#label", "Label Value"), # Should be filtered + ("entity3", "predicate3", "object3") + ] + query.get_subgraph = AsyncMock(return_value=test_subgraph) + + # Mock maybe_label to return human-readable labels + async def mock_maybe_label(entity): + label_map = { + "entity1": "Human Entity One", + "predicate1": "Human Predicate One", + "object1": "Human Object One", + "entity3": "Human Entity Three", + "predicate3": "Human Predicate Three", + "object3": "Human Object Three" + } + return label_map.get(entity, entity) + + query.maybe_label = AsyncMock(side_effect=mock_maybe_label) + + # Call get_labelgraph + result = await query.get_labelgraph("test query") + + # Verify get_subgraph was called + query.get_subgraph.assert_called_once_with("test query") + + # Verify label triples are filtered out + assert len(result) == 2 # Label triple should be excluded + + # Verify maybe_label was called for non-label triples + expected_calls = [ + (("entity1",), {}), (("predicate1",), {}), (("object1",), {}), + (("entity3",), {}), (("predicate3",), {}), (("object3",), {}) + ] + assert query.maybe_label.call_count == 6 + + # Verify result contains human-readable labels + expected_result = [ + ("Human Entity One", "Human Predicate One", "Human Object One"), + ("Human Entity Three", "Human Predicate Three", "Human Object Three") + ] + assert result == expected_result + + @pytest.mark.asyncio + async def test_graph_rag_query_method(self): + """Test GraphRag.query method orchestrates full RAG pipeline""" + # Create mock clients + mock_prompt_client = AsyncMock() + mock_embeddings_client = AsyncMock() + mock_graph_embeddings_client = AsyncMock() + mock_triples_client = AsyncMock() + + # Mock prompt client response + expected_response = "This is the RAG response" + mock_prompt_client.kg_prompt.return_value = expected_response + + # 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, + verbose=False + ) + + # Mock the Query class behavior by patching get_labelgraph + test_labelgraph = [("Subject", "Predicate", "Object")] + + # We need to patch the Query class's get_labelgraph method + original_query_init = Query.__init__ + original_get_labelgraph = Query.get_labelgraph + + def mock_query_init(self, *args, **kwargs): + original_query_init(self, *args, **kwargs) + + async def mock_get_labelgraph(self, query_text): + return test_labelgraph + + Query.__init__ = mock_query_init + Query.get_labelgraph = mock_get_labelgraph + + try: + # Call GraphRag.query + result = await graph_rag.query( + query="test query", + user="test_user", + collection="test_collection", + entity_limit=25, + triple_limit=15 + ) + + # Verify prompt client was called with knowledge graph and query + mock_prompt_client.kg_prompt.assert_called_once_with("test query", test_labelgraph) + + # Verify result + assert result == expected_response + + finally: + # Restore original methods + Query.__init__ = original_query_init + Query.get_labelgraph = original_get_labelgraph \ No newline at end of file