167 lines
5.8 KiB
Python
167 lines
5.8 KiB
Python
"""
|
||
Test script for mattergen_gen/material_gen_tools.py
|
||
|
||
This script tests the generate_material function from the material_gen_tools module.
|
||
"""
|
||
|
||
import sys
|
||
import asyncio
|
||
import unittest
|
||
import json
|
||
import re
|
||
from pathlib import Path
|
||
|
||
# 添加项目根目录到Python路径
|
||
sys.path.append(str(Path(__file__).resolve().parents[2]))
|
||
|
||
from sci_mcp.material_mcp.mattergen_gen.material_gen_tools import generate_material_MatterGen
|
||
|
||
|
||
class TestMatterGen(unittest.TestCase):
|
||
"""Test cases for MatterGen material generation tools."""
|
||
|
||
def test_unconditional_generation(self):
|
||
"""Test unconditional crystal structure generation."""
|
||
# 无条件生成（不指定属性）
|
||
result = generate_material_MatterGen(properties=None, batch_size=1, num_batches=1)
|
||
|
||
# 验证结果是否包含预期的关键信息
|
||
self.assertIsInstance(result, str)
|
||
# 检查结果是否包含一些常见的描述性文本
|
||
self.assertIn("Material", result)
|
||
self.assertIn("structures", result)
|
||
|
||
print("无条件生成结果示例:")
|
||
print(result[:500] + "...\n" if len(result) > 500 else result)
|
||
|
||
return result
|
||
|
||
# def test_single_property_generation(self):
|
||
# """Test crystal structure generation with a single property constraint."""
|
||
# # 单属性条件生成 - 使用化学系统属性
|
||
# properties = {"chemical_system": "Si-O"}
|
||
# result = generate_material(properties=properties, batch_size=1, num_batches=1)
|
||
|
||
# # 验证结果是否包含预期的关键信息
|
||
# self.assertIsInstance(result, str)
|
||
# # 检查结果是否包含相关的化学元素
|
||
# self.assertIn("Si-O", result)
|
||
|
||
# print("单属性条件生成结果示例:")
|
||
# print(result[:500] + "...\n" if len(result) > 500 else result)
|
||
|
||
# return result
|
||
|
||
# def test_multi_property_generation(self):
|
||
# """Test crystal structure generation with multiple property constraints."""
|
||
# # 多属性条件生成
|
||
# properties = {
|
||
# "chemical_system": "Fe-O",
|
||
# "space_group": 227 # 立方晶系，空间群Fd-3m
|
||
# }
|
||
# result = generate_material(properties=properties, batch_size=1, num_batches=1)
|
||
|
||
# # 验证结果是否为字符串
|
||
# self.assertIsInstance(result, str)
|
||
|
||
# # 检查结果 - 可能是成功生成或错误信息
|
||
# if "Error" in result:
|
||
# # 如果是错误信息，验证它包含相关的属性信息
|
||
# self.assertIn("properties", result)
|
||
# print("多属性条件生成返回错误 (这是预期的，因为可能不支持多属性):")
|
||
# else:
|
||
# # 如果成功，验证包含相关元素
|
||
# self.assertIn("Fe", result)
|
||
# self.assertIn("O", result)
|
||
# print("多属性条件生成成功:")
|
||
|
||
# print(result[:500] + "...\n" if len(result) > 500 else result)
|
||
|
||
# return result
|
||
|
||
# def test_batch_generation(self):
|
||
# """Test generating multiple structures in batches."""
|
||
# # 测试批量生成
|
||
# result = generate_material(properties=None, batch_size=2, num_batches=2)
|
||
|
||
# # 验证结果是否包含预期的关键信息
|
||
# self.assertIsInstance(result, str)
|
||
|
||
# # 检查结果是否提到了批量生成
|
||
# self.assertIn("structures", result)
|
||
|
||
# print("批量生成结果示例:")
|
||
# print(result[:500] + "...\n" if len(result) > 500 else result)
|
||
|
||
# return result
|
||
|
||
# def test_guidance_factor(self):
|
||
# """Test the effect of diffusion guidance factor."""
|
||
# # 测试不同的diffusion_guidance_factor值
|
||
# properties = {"chemical_system": "Al-O"}
|
||
|
||
# # 使用较低的指导因子
|
||
# result_low = generate_material(
|
||
# properties=properties,
|
||
# batch_size=1,
|
||
# num_batches=1,
|
||
# diffusion_guidance_factor=1.0
|
||
# )
|
||
|
||
# # 使用较高的指导因子
|
||
# result_high = generate_material(
|
||
# properties=properties,
|
||
# batch_size=1,
|
||
# num_batches=1,
|
||
# diffusion_guidance_factor=3.0
|
||
# )
|
||
|
||
# # 验证两个结果都是有效的
|
||
# self.assertIsInstance(result_low, str)
|
||
# self.assertIsInstance(result_high, str)
|
||
# self.assertIn("Al-O", result_low)
|
||
# self.assertIn("Al-O", result_high)
|
||
|
||
# # 验证两个结果都提到了diffusion guidance factor
|
||
# self.assertIn("guidance factor", result_low)
|
||
# self.assertIn("guidance factor", result_high)
|
||
|
||
# print("不同指导因子的生成结果示例:")
|
||
# print("低指导因子 (1.0):")
|
||
# print(result_low[:300] + "...\n" if len(result_low) > 300 else result_low)
|
||
# print("高指导因子 (3.0):")
|
||
# print(result_high[:300] + "...\n" if len(result_high) > 300 else result_high)
|
||
|
||
# return result_low, result_high
|
||
|
||
# def test_invalid_properties(self):
|
||
# """Test handling of invalid properties."""
|
||
# # 测试无效属性
|
||
# invalid_properties = {"invalid_property": "value"}
|
||
# result = generate_material(properties=invalid_properties)
|
||
|
||
# # 验证结果是否为字符串
|
||
# self.assertIsInstance(result, str)
|
||
|
||
# # 检查结果 - 可能返回错误信息或尝试生成
|
||
# if "Error" in result:
|
||
# print("无效属性测试返回错误 (预期行为):")
|
||
# else:
|
||
# # 如果没有返回错误，至少应该包含我们请求的属性名称
|
||
# self.assertIn("invalid_property", result)
|
||
# print("无效属性测试尝试生成:")
|
||
|
||
# print(result)
|
||
|
||
# return result
|
||
|
||
|
||
def run_tests():
|
||
"""运行所有测试。"""
|
||
unittest.main()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
run_tests()
|
||
|