Refactoring - fix bug