Fixed failing unit tests after refactor.