fixing failing tests after refactor