Fixing unit test failures