Fix a bugs for UnitTests