Update code for Rubocop