Fix unit tests after refactoring