Fix for unit test failures