Fix after refactor