From 0c47d1a2d1c53d5ede53d2eab14a9ccd6ddca1a9 Mon Sep 17 00:00:00 2001 From: yangdx Date: Thu, 6 Nov 2025 19:12:40 +0800 Subject: [PATCH] Fix linting --- tests/test_graph_storage.py | 242 ++++++++++++++++++++++++++---------- 1 file changed, 178 insertions(+), 64 deletions(-) diff --git a/tests/test_graph_storage.py b/tests/test_graph_storage.py index f707974b..96674046 100644 --- a/tests/test_graph_storage.py +++ b/tests/test_graph_storage.py @@ -105,7 +105,9 @@ async def initialize_graph_storage(): "vector_db_storage_cls_kwargs": { "cosine_better_than_threshold": 0.5 # Cosine similarity threshold }, - "working_dir": os.environ.get("WORKING_DIR", "./rag_storage"), # Working directory + "working_dir": os.environ.get( + "WORKING_DIR", "./rag_storage" + ), # Working directory } # If using NetworkXStorage, initialize shared_storage first @@ -173,7 +175,9 @@ async def test_graph_basic(storage): node1_props = await storage.get_node(node1_id) if node1_props: print(f"Successfully read node properties: {node1_id}") - print(f"Node description: {node1_props.get('description', 'No description')}") + print( + f"Node description: {node1_props.get('description', 'No description')}" + ) print(f"Node type: {node1_props.get('entity_type', 'No type')}") print(f"Node keywords: {node1_props.get('keywords', 'No keywords')}") # Verify that the returned properties are correct @@ -195,8 +199,12 @@ async def test_graph_basic(storage): edge_props = await storage.get_edge(node1_id, node2_id) if edge_props: print(f"Successfully read edge properties: {node1_id} -> {node2_id}") - print(f"Edge relationship: {edge_props.get('relationship', 'No relationship')}") - print(f"Edge description: {edge_props.get('description', 'No description')}") + print( + f"Edge relationship: {edge_props.get('relationship', 'No relationship')}" + ) + print( + f"Edge description: {edge_props.get('description', 'No description')}" + ) print(f"Edge weight: {edge_props.get('weight', 'No weight')}") # Verify that the returned properties are correct assert ( @@ -205,7 +213,9 @@ async def test_graph_basic(storage): assert ( edge_props.get("description") == edge_data["description"] ), "Edge description mismatch" - assert edge_props.get("weight") == edge_data["weight"], "Edge weight mismatch" + assert ( + edge_props.get("weight") == edge_data["weight"] + ), "Edge weight mismatch" else: print(f"Failed to read edge properties: {node1_id} -> {node2_id}") assert False, f"Failed to read edge properties: {node1_id} -> {node2_id}" @@ -214,20 +224,28 @@ async def test_graph_basic(storage): print(f"Reading reverse edge properties: {node2_id} -> {node1_id}") reverse_edge_props = await storage.get_edge(node2_id, node1_id) if reverse_edge_props: - print(f"Successfully read reverse edge properties: {node2_id} -> {node1_id}") - print(f"Reverse edge relationship: {reverse_edge_props.get('relationship', 'No relationship')}") - print(f"Reverse edge description: {reverse_edge_props.get('description', 'No description')}") - print(f"Reverse edge weight: {reverse_edge_props.get('weight', 'No weight')}") + print( + f"Successfully read reverse edge properties: {node2_id} -> {node1_id}" + ) + print( + f"Reverse edge relationship: {reverse_edge_props.get('relationship', 'No relationship')}" + ) + print( + f"Reverse edge description: {reverse_edge_props.get('description', 'No description')}" + ) + print( + f"Reverse edge weight: {reverse_edge_props.get('weight', 'No weight')}" + ) # Verify that forward and reverse edge properties are the same assert ( edge_props == reverse_edge_props ), "Forward and reverse edge properties are not consistent, undirected graph property verification failed" - print("Undirected graph property verification successful: forward and reverse edge properties are consistent") + print( + "Undirected graph property verification successful: forward and reverse edge properties are consistent" + ) else: print(f"Failed to read reverse edge properties: {node2_id} -> {node1_id}") - assert ( - False - ), f"Failed to read reverse edge properties: {node2_id} -> {node1_id}, undirected graph property verification failed" + assert False, f"Failed to read reverse edge properties: {node2_id} -> {node1_id}, undirected graph property verification failed" print("Basic tests completed, data is preserved in the database.") return True @@ -307,7 +325,9 @@ async def test_graph_advanced(storage): print(f"== Testing node_degree: {node1_id}") node1_degree = await storage.node_degree(node1_id) print(f"Degree of node {node1_id}: {node1_degree}") - assert node1_degree == 1, f"Degree of node {node1_id} should be 1, but got {node1_degree}" + assert ( + node1_degree == 1 + ), f"Degree of node {node1_id} should be 1, but got {node1_degree}" # 2.1 Test degrees of all nodes print("== Testing degrees of all nodes") @@ -315,8 +335,12 @@ async def test_graph_advanced(storage): node3_degree = await storage.node_degree(node3_id) print(f"Degree of node {node2_id}: {node2_degree}") print(f"Degree of node {node3_id}: {node3_degree}") - assert node2_degree == 2, f"Degree of node {node2_id} should be 2, but got {node2_degree}" - assert node3_degree == 1, f"Degree of node {node3_id} should be 1, but got {node3_degree}" + assert ( + node2_degree == 2 + ), f"Degree of node {node2_id} should be 2, but got {node2_degree}" + assert ( + node3_degree == 1 + ), f"Degree of node {node3_id} should be 1, but got {node3_degree}" # 3. Test edge_degree - get the degree of an edge print(f"== Testing edge_degree: {node1_id} -> {node2_id}") @@ -333,7 +357,9 @@ async def test_graph_advanced(storage): assert ( edge_degree == reverse_edge_degree ), "Degrees of forward and reverse edges are not consistent, undirected graph property verification failed" - print("Undirected graph property verification successful: degrees of forward and reverse edges are consistent") + print( + "Undirected graph property verification successful: degrees of forward and reverse edges are consistent" + ) # 4. Test get_node_edges - get all edges of a node print(f"== Testing get_node_edges: {node2_id}") @@ -366,7 +392,9 @@ async def test_graph_advanced(storage): assert ( has_connection_with_node3 ), f"Edge list of node {node2_id} should include a connection with {node3_id}" - print(f"Undirected graph property verification successful: edge list of node {node2_id} contains all relevant edges") + print( + f"Undirected graph property verification successful: edge list of node {node2_id} contains all relevant edges" + ) # 5. Test get_all_labels - get all labels print("== Testing get_all_labels") @@ -382,9 +410,15 @@ async def test_graph_advanced(storage): kg = await storage.get_knowledge_graph("*", max_depth=2, max_nodes=10) print(f"Number of nodes in knowledge graph: {len(kg.nodes)}") print(f"Number of edges in knowledge graph: {len(kg.edges)}") - assert isinstance(kg, KnowledgeGraph), "The returned result should be of type KnowledgeGraph" - assert len(kg.nodes) == 3, f"The knowledge graph should have 3 nodes, but got {len(kg.nodes)}" - assert len(kg.edges) == 2, f"The knowledge graph should have 2 edges, but got {len(kg.edges)}" + assert isinstance( + kg, KnowledgeGraph + ), "The returned result should be of type KnowledgeGraph" + assert ( + len(kg.nodes) == 3 + ), f"The knowledge graph should have 3 nodes, but got {len(kg.nodes)}" + assert ( + len(kg.edges) == 2 + ), f"The knowledge graph should have 2 edges, but got {len(kg.edges)}" # 7. Test delete_node - delete a node print(f"== Testing delete_node: {node3_id}") @@ -401,17 +435,27 @@ async def test_graph_advanced(storage): print(f"== Testing remove_edges: {node2_id} -> {node3_id}") await storage.remove_edges([(node2_id, node3_id)]) edge_props = await storage.get_edge(node2_id, node3_id) - print(f"Querying edge properties after deletion {node2_id} -> {node3_id}: {edge_props}") - assert edge_props is None, f"Edge {node2_id} -> {node3_id} should have been deleted" + print( + f"Querying edge properties after deletion {node2_id} -> {node3_id}: {edge_props}" + ) + assert ( + edge_props is None + ), f"Edge {node2_id} -> {node3_id} should have been deleted" # 8.1 Verify undirected graph property of edge deletion - print(f"== Verifying undirected graph property of edge deletion: {node3_id} -> {node2_id}") + print( + f"== Verifying undirected graph property of edge deletion: {node3_id} -> {node2_id}" + ) reverse_edge_props = await storage.get_edge(node3_id, node2_id) - print(f"Querying reverse edge properties after deletion {node3_id} -> {node2_id}: {reverse_edge_props}") + print( + f"Querying reverse edge properties after deletion {node3_id} -> {node2_id}: {reverse_edge_props}" + ) assert ( reverse_edge_props is None ), f"Reverse edge {node3_id} -> {node2_id} should also be deleted, undirected graph property verification failed" - print("Undirected graph property verification successful: deleting an edge in one direction also deletes the reverse edge") + print( + "Undirected graph property verification successful: deleting an edge in one direction also deletes the reverse edge" + ) # 9. Test remove_nodes - delete multiple nodes print(f"== Testing remove_nodes: [{node2_id}, {node3_id}]") @@ -636,7 +680,9 @@ async def test_graph_batch_operations(storage): edge_dicts = [{"src": src, "tgt": tgt} for src, tgt in edges] edges_dict = await storage.get_edges_batch(edge_dicts) print(f"Batch get edge properties result: {edges_dict.keys()}") - assert len(edges_dict) == 3, f"Should return properties of 3 edges, but got {len(edges_dict)}" + assert ( + len(edges_dict) == 3 + ), f"Should return properties of 3 edges, but got {len(edges_dict)}" assert ( node1_id, node2_id, @@ -675,14 +721,19 @@ async def test_graph_batch_operations(storage): # Verify that properties of forward and reverse edges are consistent for (src, tgt), props in edges_dict.items(): assert ( - tgt, - src, - ) in reverse_edges_dict, f"Reverse edge {tgt} -> {src} should be in the result" + ( + tgt, + src, + ) + in reverse_edges_dict + ), f"Reverse edge {tgt} -> {src} should be in the result" assert ( props == reverse_edges_dict[(tgt, src)] ), f"Properties of edge {src} -> {tgt} and reverse edge {tgt} -> {src} are inconsistent" - print("Undirected graph property verification successful: properties of batch-retrieved forward and reverse edges are consistent") + print( + "Undirected graph property verification successful: properties of batch-retrieved forward and reverse edges are consistent" + ) # 6. Test get_nodes_edges_batch - batch get all edges of multiple nodes print("== Testing get_nodes_edges_batch") @@ -718,9 +769,15 @@ async def test_graph_batch_operations(storage): has_edge_to_node4 = any(tgt == node4_id for _, tgt in node1_outgoing_edges) has_edge_to_node5 = any(tgt == node5_id for _, tgt in node1_outgoing_edges) - assert has_edge_to_node2, f"Edge list of node {node1_id} should include an edge to {node2_id}" - assert has_edge_to_node4, f"Edge list of node {node1_id} should include an edge to {node4_id}" - assert has_edge_to_node5, f"Edge list of node {node1_id} should include an edge to {node5_id}" + assert ( + has_edge_to_node2 + ), f"Edge list of node {node1_id} should include an edge to {node2_id}" + assert ( + has_edge_to_node4 + ), f"Edge list of node {node1_id} should include an edge to {node4_id}" + assert ( + has_edge_to_node5 + ), f"Edge list of node {node1_id} should include an edge to {node5_id}" # Check if node 3's edges include all relevant edges (regardless of direction) node3_outgoing_edges = [ @@ -759,7 +816,9 @@ async def test_graph_batch_operations(storage): has_connection_with_node5 ), f"Edge list of node {node3_id} should include a connection with {node5_id}" - print("Undirected graph property verification successful: batch-retrieved node edges include all relevant edges (regardless of direction)") + print( + "Undirected graph property verification successful: batch-retrieved node edges include all relevant edges (regardless of direction)" + ) print("\nBatch operations tests completed.") return True @@ -825,7 +884,9 @@ async def test_graph_special_characters(storage): "weight": 0.8, "description": "Contains SQL injection attempt: SELECT * FROM users WHERE name='admin'--", } - print(f"Inserting edge with complex special characters: {node2_id} -> {node3_id}") + print( + f"Inserting edge with complex special characters: {node2_id} -> {node3_id}" + ) await storage.upsert_edge(node2_id, node3_id, edge2_data) # 6. Verify that node special characters are saved correctly @@ -838,7 +899,9 @@ async def test_graph_special_characters(storage): node_props = await storage.get_node(node_id) if node_props: print(f"Successfully read node: {node_id}") - print(f"Node description: {node_props.get('description', 'No description')}") + print( + f"Node description: {node_props.get('description', 'No description')}" + ) # Verify node ID is saved correctly assert ( @@ -860,8 +923,12 @@ async def test_graph_special_characters(storage): edge1_props = await storage.get_edge(node1_id, node2_id) if edge1_props: print(f"Successfully read edge: {node1_id} -> {node2_id}") - print(f"Edge relationship: {edge1_props.get('relationship', 'No relationship')}") - print(f"Edge description: {edge1_props.get('description', 'No description')}") + print( + f"Edge relationship: {edge1_props.get('relationship', 'No relationship')}" + ) + print( + f"Edge description: {edge1_props.get('description', 'No description')}" + ) # Verify edge relationship is saved correctly assert ( @@ -873,7 +940,9 @@ async def test_graph_special_characters(storage): edge1_props.get("description") == edge1_data["description"] ), f"Edge description mismatch: expected {edge1_data['description']}, got {edge1_props.get('description')}" - print(f"Edge {node1_id} -> {node2_id} special character verification successful") + print( + f"Edge {node1_id} -> {node2_id} special character verification successful" + ) else: print(f"Failed to read edge properties: {node1_id} -> {node2_id}") assert False, f"Failed to read edge properties: {node1_id} -> {node2_id}" @@ -881,8 +950,12 @@ async def test_graph_special_characters(storage): edge2_props = await storage.get_edge(node2_id, node3_id) if edge2_props: print(f"Successfully read edge: {node2_id} -> {node3_id}") - print(f"Edge relationship: {edge2_props.get('relationship', 'No relationship')}") - print(f"Edge description: {edge2_props.get('description', 'No description')}") + print( + f"Edge relationship: {edge2_props.get('relationship', 'No relationship')}" + ) + print( + f"Edge description: {edge2_props.get('description', 'No description')}" + ) # Verify edge relationship is saved correctly assert ( @@ -894,7 +967,9 @@ async def test_graph_special_characters(storage): edge2_props.get("description") == edge2_data["description"] ), f"Edge description mismatch: expected {edge2_data['description']}, got {edge2_props.get('description')}" - print(f"Edge {node2_id} -> {node3_id} special character verification successful") + print( + f"Edge {node2_id} -> {node3_id} special character verification successful" + ) else: print(f"Failed to read edge properties: {node2_id} -> {node3_id}") assert False, f"Failed to read edge properties: {node2_id} -> {node3_id}" @@ -965,18 +1040,24 @@ async def test_graph_undirected_property(storage): # Verify forward query forward_edge = await storage.get_edge(node1_id, node2_id) print(f"Forward edge properties: {forward_edge}") - assert forward_edge is not None, f"Failed to read forward edge properties: {node1_id} -> {node2_id}" + assert ( + forward_edge is not None + ), f"Failed to read forward edge properties: {node1_id} -> {node2_id}" # Verify reverse query reverse_edge = await storage.get_edge(node2_id, node1_id) print(f"Reverse edge properties: {reverse_edge}") - assert reverse_edge is not None, f"Failed to read reverse edge properties: {node2_id} -> {node1_id}" + assert ( + reverse_edge is not None + ), f"Failed to read reverse edge properties: {node2_id} -> {node1_id}" # Verify that forward and reverse edge properties are consistent assert ( forward_edge == reverse_edge ), "Forward and reverse edge properties are inconsistent, undirected property verification failed" - print("Undirected property verification successful: forward and reverse edge properties are consistent") + print( + "Undirected property verification successful: forward and reverse edge properties are consistent" + ) # 3. Test undirected property of edge degree print("\n== Testing undirected property of edge degree") @@ -998,7 +1079,9 @@ async def test_graph_undirected_property(storage): assert ( forward_degree == reverse_degree ), "Degrees of forward and reverse edges are inconsistent, undirected property verification failed" - print("Undirected property verification successful: degrees of forward and reverse edges are consistent") + print( + "Undirected property verification successful: degrees of forward and reverse edges are consistent" + ) # 4. Test undirected property of edge deletion print("\n== Testing undirected property of edge deletion") @@ -1009,16 +1092,24 @@ async def test_graph_undirected_property(storage): # Verify forward edge is deleted forward_edge = await storage.get_edge(node1_id, node2_id) - print(f"Querying forward edge properties after deletion {node1_id} -> {node2_id}: {forward_edge}") - assert forward_edge is None, f"Edge {node1_id} -> {node2_id} should have been deleted" + print( + f"Querying forward edge properties after deletion {node1_id} -> {node2_id}: {forward_edge}" + ) + assert ( + forward_edge is None + ), f"Edge {node1_id} -> {node2_id} should have been deleted" # Verify reverse edge is also deleted reverse_edge = await storage.get_edge(node2_id, node1_id) - print(f"Querying reverse edge properties after deletion {node2_id} -> {node1_id}: {reverse_edge}") + print( + f"Querying reverse edge properties after deletion {node2_id} -> {node1_id}: {reverse_edge}" + ) assert ( reverse_edge is None ), f"Reverse edge {node2_id} -> {node1_id} should also be deleted, undirected property verification failed" - print("Undirected property verification successful: deleting an edge in one direction also deletes the reverse edge") + print( + "Undirected property verification successful: deleting an edge in one direction also deletes the reverse edge" + ) # 5. Test undirected property in batch operations print("\n== Testing undirected property in batch operations") @@ -1045,14 +1136,19 @@ async def test_graph_undirected_property(storage): # Verify that properties of forward and reverse edges are consistent for (src, tgt), props in edges_dict.items(): assert ( - tgt, - src, - ) in reverse_edges_dict, f"Reverse edge {tgt} -> {src} should be in the result" + ( + tgt, + src, + ) + in reverse_edges_dict + ), f"Reverse edge {tgt} -> {src} should be in the result" assert ( props == reverse_edges_dict[(tgt, src)] ), f"Properties of edge {src} -> {tgt} and reverse edge {tgt} -> {src} are inconsistent" - print("Undirected property verification successful: properties of batch-retrieved forward and reverse edges are consistent") + print( + "Undirected property verification successful: properties of batch-retrieved forward and reverse edges are consistent" + ) # 6. Test undirected property of batch-retrieved node edges print("\n== Testing undirected property of batch-retrieved node edges") @@ -1072,8 +1168,12 @@ async def test_graph_undirected_property(storage): (src == node1_id and tgt == node3_id) for src, tgt in node1_edges ) - assert has_edge_to_node2, f"Edge list of node {node1_id} should include an edge to {node2_id}" - assert has_edge_to_node3, f"Edge list of node {node1_id} should include an edge to {node3_id}" + assert ( + has_edge_to_node2 + ), f"Edge list of node {node1_id} should include an edge to {node2_id}" + assert ( + has_edge_to_node3 + ), f"Edge list of node {node1_id} should include an edge to {node3_id}" # Check if node 2 has a connection with node 1 has_edge_to_node1 = any( @@ -1085,7 +1185,9 @@ async def test_graph_undirected_property(storage): has_edge_to_node1 ), f"Edge list of node {node2_id} should include a connection with {node1_id}" - print("Undirected property verification successful: batch-retrieved node edges include all relevant edges (regardless of direction)") + print( + "Undirected property verification successful: batch-retrieved node edges include all relevant edges (regardless of direction)" + ) print("\nUndirected property tests completed.") return True @@ -1113,7 +1215,9 @@ async def main(): # Get graph storage type graph_storage_type = os.getenv("LIGHTRAG_GRAPH_STORAGE", "NetworkXStorage") - ASCIIColors.magenta(f"\nCurrently configured graph storage type: {graph_storage_type}") + ASCIIColors.magenta( + f"\nCurrently configured graph storage type: {graph_storage_type}" + ) ASCIIColors.white( f"Supported graph storage types: {', '.join(STORAGE_IMPLEMENTATIONS['GRAPH_STORAGE']['implementations'])}" ) @@ -1128,10 +1232,18 @@ async def main(): # Display test options ASCIIColors.yellow("\nPlease select a test type:") ASCIIColors.white("1. Basic Test (Node and edge insertion, reading)") - ASCIIColors.white("2. Advanced Test (Degree, labels, knowledge graph, deletion, etc.)") - ASCIIColors.white("3. Batch Operations Test (Batch get node/edge properties, degrees, etc.)") - ASCIIColors.white("4. Undirected Property Test (Verify undirected properties of the storage)") - ASCIIColors.white("5. Special Characters Test (Verify handling of single/double quotes, backslashes, etc.)") + ASCIIColors.white( + "2. Advanced Test (Degree, labels, knowledge graph, deletion, etc.)" + ) + ASCIIColors.white( + "3. Batch Operations Test (Batch get node/edge properties, degrees, etc.)" + ) + ASCIIColors.white( + "4. Undirected Property Test (Verify undirected properties of the storage)" + ) + ASCIIColors.white( + "5. Special Characters Test (Verify handling of single/double quotes, backslashes, etc.)" + ) ASCIIColors.white("6. All Tests") choice = input("\nEnter your choice (1/2/3/4/5/6): ") @@ -1171,7 +1283,9 @@ async def main(): ) if undirected_result: - ASCIIColors.cyan("\n=== Starting Special Characters Test ===") + ASCIIColors.cyan( + "\n=== Starting Special Characters Test ===" + ) await test_graph_special_characters(storage) else: ASCIIColors.red("Invalid choice")