Fix test failures from refactoring