Fixing bug in refactored code (missing variable)