Fix some unit test