增加单元测试和优化代码