knowledge/tests/community/test_dump_graph_data.py

import unittest
import unittest
import json
import os
from community.dump_graph_data import get_props, get_entities, get_relationships

class TestDumpGraphData(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.test_data_path = os.path.join(os.getcwd(), 'web', 'cached_data')
        os.makedirs(cls.test_data_path, exist_ok=True)

    def test_get_props(self):
        """测试属性获取方法，验证多属性合并逻辑"""
        props = get_props('1')  # 使用真实节点ID
        self.assertIsInstance(props, dict)
        self.assertTrue(len(props) > 0)

    def test_entity_export(self):
        """验证实体导出的分页查询和JSON序列化"""
        get_entities()
        
        # 验证生成的实体文件
        with open(os.path.join(self.test_data_path, 'entities_med.json'), 'r', encoding='utf-8') as f:
            entities = json.load(f)
            self.assertTrue(len(entities) > 0)

    def test_relationship_chunking(self):
        """测试关系数据分块写入逻辑，验证每批处理1000条"""
        get_relationships()
        
        # 验证生成的关系文件
        file_count = len([f for f in os.listdir(self.test_data_path) 
                         if f.startswith('relationship_med_')])
        self.assertTrue(file_count > 0)

    #def tearDown(self):
        # 清理测试生成的文件
        #for f in os.listdir(self.test_data_path):
            #if f.startswith(('entities_med', 'relationship_med_')):
                #os.remove(os.path.join(self.test_data_path, f))

if __name__ == '__main__':
    unittest.main()