Fix unit test bugs