LightRAG源码：NetworkXStorage测试(1)

测试代码

nano-graphrag跟lightrag代码雷同，测试案例来自与nano-graphrag

1. import os
2. import shutil
3. import pytest
4. import networkx as nx
5. import numpy as np
6. import asyncio
8. import json
9. from lightrag import LightRAG
11. from lightrag.storage import NetworkXStorage
12. from lightrag.utils import wrap_embedding_func_with_attrs
14. WORKING_DIR = "./tests/nano_graphrag_cache_networkx_storage_test"
17. @pytest.fixture(scope="function")
18. def setup_teardown():
19.     if os.path.exists(WORKING_DIR):
20.         shutil.rmtree(WORKING_DIR)
21.     os.mkdir(WORKING_DIR)
23.     yield
25.     shutil.rmtree(WORKING_DIR)
28. @wrap_embedding_func_with_attrs(embedding_dim=384, max_token_size=8192)
29. async def mock_embedding(texts: list[str]) -> np.ndarray:
30.     return np.random.rand(len(texts), 384)
33. @pytest.fixture
34. def networkx_storage(setup_teardown):
35.     rag = LightRAG(working_dir=WORKING_DIR, embedding_func=mock_embedding)
36.     return NetworkXStorage(
37.         namespace="test",
38.         global_config=rag.__dict__,
39.     )
42. @pytest.mark.asyncio
43. async def test_upsert_and_get_node(networkx_storage):
44.     node_id = "node1"
45.     node_data = {"attr1": "value1", "attr2": "value2"}
47.     await networkx_storage.upsert_node(node_id, node_data)
49.     result = await networkx_storage.get_node(node_id)
50.     assert result == node_data
52.     has_node = await networkx_storage.has_node(node_id)
53.     assert has_node is True
56. @pytest.mark.asyncio
57. async def test_upsert_and_get_edge(networkx_storage):
58.     source_id = "node1"
59.     target_id = "node2"
60.     edge_data = {"weight": 1.0, "type": "connection"}
62.     await networkx_storage.upsert_node(source_id, {})
63.     await networkx_storage.upsert_node(target_id, {})
64.     await networkx_storage.upsert_edge(source_id, target_id, edge_data)
66.     result = await networkx_storage.get_edge(source_id, target_id)
67.     assert result == edge_data
69.     has_edge = await networkx_storage.has_edge(source_id, target_id)
70.     assert has_edge is True
73. @pytest.mark.asyncio
74. async def test_node_degree(networkx_storage):
75.     node_id = "center"
76.     await networkx_storage.upsert_node(node_id, {})
78.     num_neighbors = 5
79.     for i in range(num_neighbors):
80.         neighbor_id = f"neighbor{i}"
81.         await networkx_storage.upsert_node(neighbor_id, {})
82.         await networkx_storage.upsert_edge(node_id, neighbor_id, {})
84.     degree = await networkx_storage.node_degree(node_id)
85.     assert degree == num_neighbors
88. @pytest.mark.asyncio
89. async def test_edge_degree(networkx_storage):
90.     source_id = "node1"
91.     target_id = "node2"
93.     await networkx_storage.upsert_node(source_id, {})
94.     await networkx_storage.upsert_node(target_id, {})
95.     await networkx_storage.upsert_edge(source_id, target_id, {})
97.     num_source_neighbors = 3
98.     for i in range(num_source_neighbors):
99.         neighbor_id = f"neighbor{i}"
100.         await networkx_storage.upsert_node(neighbor_id, {})
101.         await networkx_storage.upsert_edge(source_id, neighbor_id, {})
103.     num_target_neighbors = 2
104.     for i in range(num_target_neighbors):
105.         neighbor_id = f"target_neighbor{i}"
106.         await networkx_storage.upsert_node(neighbor_id, {})
107.         await networkx_storage.upsert_edge(target_id, neighbor_id, {})
109.     expected_edge_degree = (num_source_neighbors + 1) + (num_target_neighbors + 1)
110.     edge_degree = await networkx_storage.edge_degree(source_id, target_id)
111.     assert edge_degree == expected_edge_degree
114. @pytest.mark.asyncio
115. async def test_get_node_edges(networkx_storage):
116.     center_id = "center"
117.     await networkx_storage.upsert_node(center_id, {})
119.     expected_edges = []
120.     for i in range(3):
121.         neighbor_id = f"neighbor{i}"
122.         await networkx_storage.upsert_node(neighbor_id, {})
123.         await networkx_storage.upsert_edge(center_id, neighbor_id, {})
124.         expected_edges.append((center_id, neighbor_id))
126.     result = await networkx_storage.get_node_edges(center_id)
127.     assert set(result) == set(expected_edges)

复制代码

代码运行逻辑解释

1. 导入依赖库

• os 和 shutil：用于文件和目录操作。
  
• pytest：用于编写和运行测试。
  
• networkx as nx：用于创建和操作图布局。
  
• numpy as np：用于生成随机数。
  
• asyncio：用于异步编程。
  
• json：用于处置惩罚JSON数据。
  
• lightrag 干系模块：用于实现RAG（Retrieval-Augmented Generation）模子的干系功能。
  

2. 界说工作目录

• WORKING_DIR：界说了一个工作目录路径，用于存储测试过程中生成的文件。
  

3. 界说 setup_teardown 夹具

• 该夹具在每个测试函数运行前后执行。
  
  
  
  • setup：在测试开始前，查抄并删除工作目录（假如存在），然后重新创建该目录。
    
  • teardown：在测试结束后，删除工作目录及其内容。
    
  
  

4. 界说 mock_embedding 函数

• 这是一个模仿的嵌入函数，用于生成随机嵌入向量。
  
  
  
  • 利用 @wrap_embedding_func_with_attrs 装饰器，指定嵌入向量的维度和最大token巨细。
    
  • 返回一个随机生成的 numpy 数组，形状为 (len(texts), 384)。
    
  
  

5. 界说 networkx_storage 夹具

• 该夹具用于创建 NetworkXStorage 实例。
  
  
  
  • 利用 LightRAG 类初始化 NetworkXStorage，并传入工作目录和模仿的嵌入函数。
    
  • 返回一个 NetworkXStorage 实例，用于后续测试。
    
  
  

6. 测试函数

6.1 test_upsert_and_get_node

• 功能：测试节点的插入和获取。
  
  
  
  • 插入一个节点 node1，并为其添加一些属性。
    
  • 利用 get_node 方法获取该节点，并验证返回的数据是否正确。
    
  • 利用 has_node 方法查抄节点是否存在。
    
  
  

6.2 test_upsert_and_get_edge

• 功能：测试边的插入和获取。
  
  
  
  • 插入两个节点 node1 和 node2。
    
  • 插入一条边，并为其添加一些属性。
    
  • 利用 get_edge 方法获取该边，并验证返回的数据是否正确。
    
  • 利用 has_edge 方法查抄边是否存在。
    
  
  

6.3 test_node_degree

• 功能：测试节点的度数（即节点的邻居数目）。
  
  
  
  • 插入一个中央节点 center。
    
  • 插入多个邻居节点，并将它们与中央节点连接。
    
  • 利用 node_degree 方法获取中央节点的度数，并验证其是否正确。
    
  
  

6.4 test_edge_degree

• 功能：测试边的度数（即与边相连的节点总数）。
  
  
  
  • 插入两个节点 node1 和 node2，并将它们连接。
    
  • 插入多个邻居节点，并将它们分别与 node1 和 node2 连接。
    
  • 利用 edge_degree 方法获取边的度数，并验证其是否正确。
    
  
  

6.5 test_get_node_edges

• 功能：测试获取节点的所有边。
  
  
  
  • 插入一个中央节点 center。
    
  • 插入多个邻居节点，并将它们与中央节点连接。
    
  • 利用 get_node_edges 方法获取中央节点的所有边，并验证返回的边是否正确。
    
  
  

7. 总结

• 这些测试函数主要用于验证 NetworkXStorage 类的功能，包括节点的插入、获取、边的插入、获取、节点和边的度数计算等。
  
• 通过 pytest 和 asyncio 的结合，实现了异步测试，确保代码在异步环境下的正确性。
  

这些测试代码的主要目的是验证 NetworkXStorage 类的功能，确保其可以或许正确地处置惩罚节点和边的插入、获取、度数计算等操作。

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。