Test fixes due to refactoring