multi_mcp/test_tools/material/test_pymatgen_cal.py

"""
Test script for pymatgen_cal_tools.py

This script tests the functions from the pymatgen_cal_tools module.
"""

import os
import sys
import asyncio
import unittest
from pathlib import Path

# 添加项目根目录到Python路径
sys.path.append(str(Path(__file__).resolve().parents[2]))

from sci_mcp.material_mcp.pymatgen_cal.pymatgen_cal_tools import (
    calculate_density,
    get_element_composition,
    calculate_symmetry
)


class TestPymatgenCalculations(unittest.TestCase):
    """Test cases for pymatgen calculation tools."""

    def setUp(self):
        """Set up test fixtures."""
        # 简单的CIF字符串示例 - 硅晶体结构
        self.simple_cif = """
data_Si
_cell_length_a 5.43
_cell_length_b 5.43
_cell_length_c 5.43
_cell_angle_alpha 90
_cell_angle_beta 90
_cell_angle_gamma 90
_symmetry_space_group_name_H-M 'P 1'
_symmetry_Int_Tables_number 1
loop_
_atom_site_label
_atom_site_fract_x
_atom_site_fract_y
_atom_site_fract_z
Si 0.0 0.0 0.0
Si 0.5 0.5 0.0
Si 0.5 0.0 0.5
Si 0.0 0.5 0.5
Si 0.25 0.25 0.25
Si 0.75 0.75 0.25
Si 0.75 0.25 0.75
Si 0.25 0.75 0.75
"""

    def test_calculate_density_async(self):
        """Test calculate_density function with a simple CIF string (异步版本)."""
        async def _async_test():
            result = await calculate_density(self.simple_cif)

            # 验证结果是否包含预期的关键信息
            self.assertIsInstance(result, str)
            self.assertIn("Density Calculation", result)
            self.assertIn("Density", result)
            self.assertIn("g/cm³", result)

            print("密度计算结果示例:")
            print(result)
            return result

        loop = asyncio.get_event_loop()
        return loop.run_until_complete(_async_test())

    def test_get_element_composition_async(self):
        """Test get_element_composition function with a simple CIF string (异步版本)."""
        async def _async_test():
            result = await get_element_composition(self.simple_cif)

            # 验证结果是否包含预期的关键信息
            self.assertIsInstance(result, str)
            self.assertIn("Element Composition", result)
            self.assertIn("Composition", result)
            self.assertIn("Si", result)

            print("元素组成计算结果示例:")
            print(result)
            return result

        loop = asyncio.get_event_loop()
        return loop.run_until_complete(_async_test())

    def test_calculate_symmetry_async(self):
        """Test calculate_symmetry function with a simple CIF string (异步版本)."""
        async def _async_test():
            result = await calculate_symmetry(self.simple_cif)

            # 验证结果是否包含预期的关键信息
            self.assertIsInstance(result, str)
            self.assertIn("Symmetry Information", result)
            self.assertIn("Space Group", result)
            self.assertIn("Number", result)

            print("对称性计算结果示例:")
            print(result)
            return result

        loop = asyncio.get_event_loop()
        return loop.run_until_complete(_async_test())


class TestPymatgenCalculationsWithFile(unittest.TestCase):
    """使用文件测试pymatgen计算工具。"""

    def setUp(self):
        """设置测试夹具，创建临时CIF文件。"""
        self.temp_cif_path = "temp_test_structure.cif"

        # 简单的CIF内容 - 氧化铝结构
        cif_content = """
data_Al2O3
_cell_length_a 4.76
_cell_length_b 4.76
_cell_length_c 12.99
_cell_angle_alpha 90
_cell_angle_beta 90
_cell_angle_gamma 120
_symmetry_space_group_name_H-M 'R -3 c'
_symmetry_Int_Tables_number 167
_symmetry_equiv_pos_as_xyz 'x, y, z'
_symmetry_equiv_pos_as_xyz '-x, -y, -z'
loop_
_atom_site_label
_atom_site_type_symbol
_atom_site_fract_x
_atom_site_fract_y
_atom_site_fract_z
_atom_site_occupancy
Al1 Al 0.0 0.0 0.35 1.0
Al2 Al 0.0 0.0 0.85 1.0
O1 O 0.31 0.0 0.25 1.0
"""

        # 创建临时文件
        with open(self.temp_cif_path, "w") as f:
            f.write(cif_content)

    def tearDown(self):
        """清理测试夹具，删除临时文件。"""
        if os.path.exists(self.temp_cif_path):
            os.remove(self.temp_cif_path)

    def test_calculate_density_from_file_async(self):
        """测试从文件计算密度（异步版本）。"""
        async def _async_test():
            result = await calculate_density(self.temp_cif_path)

            # 验证结果
            self.assertIsInstance(result, str)
            self.assertIn("Density Calculation", result)
            self.assertIn("Density", result)

            print("从文件计算密度结果示例:")
            print(result)
            return result

        loop = asyncio.get_event_loop()
        return loop.run_until_complete(_async_test())

    def test_get_element_composition_from_file_async(self):
        """测试从文件获取元素组成（异步版本）。"""
        async def _async_test():
            result = await get_element_composition(self.temp_cif_path)

            # 验证结果
            self.assertIsInstance(result, str)
            self.assertIn("Element Composition", result)
            self.assertIn("Composition", result)

            print("从文件获取元素组成结果示例:")
            print(result)
            return result

        loop = asyncio.get_event_loop()
        return loop.run_until_complete(_async_test())

    def test_calculate_symmetry_from_file_async(self):
        """测试从文件计算对称性（异步版本）。"""
        async def _async_test():
            result = await calculate_symmetry(self.temp_cif_path)

            # 验证结果
            self.assertIsInstance(result, str)
            self.assertIn("Symmetry Information", result)
            self.assertIn("Space Group", result)

            print("从文件计算对称性结果示例:")
            print(result)
            return result

        loop = asyncio.get_event_loop()
        return loop.run_until_complete(_async_test())


def run_tests():
    """运行所有测试。"""
    unittest.main()


if __name__ == "__main__":
    run_tests()