整理代码，加入单元测试功能