Fixing unit test bugs