Unit test refactoring fixes