Refactor bug fix