Fixing some unit test failures