Fix wrong autocorrect for Rails/FindByOrAssignmentMemoization

Followup to 501e5d3

Author: Earlopain

changelog/fix_autocorrect_for_find_by_memo.md

@@ -0,0 +1 @@
+* [#1516](https://github.com/rubocop/rubocop-rails/pull/1516): Fix wrong autocorrect for `Rails/FindByOrAssignmentMemoization`. ([@earlopain][])

lib/rubocop/cop/rails/find_by_or_assignment_memoization.rb

@@ -13,7 +13,7 @@ module Rails
  # or the code may have a different purpose than memoization.
  #
  # @example
- # # bad
+ # # bad - exclusively doing memoization
  # def current_user
  # @current_user ||= User.find_by(id: session[:user_id])
  # end
@@ -24,6 +24,22 @@ module Rails
  #
  # @current_user = User.find_by(id: session[:user_id])
  # end
+ #
+ # # bad - method contains other code
+ # def current_user
+ # @current_user ||= User.find_by(id: session[:user_id])
+ # @current_user.do_something
+ # end
+ #
+ # # good
+ # def current_user
+ # if defined?(@current_user)
+ # @current_user
+ # else
+ # @current_user = User.find_by(id: session[:user_id])
+ # end
+ # @current_user.do_something
+ # end
  class FindByOrAssignmentMemoization < Base
  extend AutoCorrector
 
@@ -38,6 +54,22 @@ class FindByOrAssignmentMemoization < Base
  )
  PATTERN
 
+ # When a method body contains only memoization, the correction can be more succinct.
+ def on_def(node)
+ find_by_or_assignment_memoization(node.body) do |varible_name, find_by|
+ add_offense(node.body) do |corrector|
+ corrector.replace(
+ node.body,
+ <<~RUBY.rstrip
+ return #{varible_name} if defined?(#{varible_name})
+ 
+ #{varible_name} = #{find_by.source}
+ RUBY
+ )
+ end
+ end
+ end
+
  def on_send(node)
  assignment_node = node.parent
  find_by_or_assignment_memoization(assignment_node) do |varible_name, find_by|
@@ -47,9 +79,11 @@ def on_send(node)
  corrector.replace(
  assignment_node,
  <<~RUBY.rstrip
- return #{varible_name} if defined?(#{varible_name})
-
- #{varible_name} = #{find_by.source}
+ if defined?(#{varible_name})
+ #{varible_name}
+ else
+ #{varible_name} = #{find_by.source}
+ end
  RUBY
  )
  end

spec/rubocop/cop/rails/find_by_or_assignment_memoization_spec.rb

@@ -9,9 +9,53 @@
  RUBY
 
  expect_correction(<<~RUBY)
- return @current_user if defined?(@current_user)
+ if defined?(@current_user)
+ @current_user
+ else
+ @current_user = User.find_by(id: session[:user_id])
+ end
+ RUBY
+ end
+ end
+
+ context 'when using `find_by` with `||=` in a method body' do
+ it 'registers an offense' do
+ expect_offense(<<~RUBY)
+ def foo
+ @current_user ||= User.find_by(id: session[:user_id])
+ ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ Avoid memoizing `find_by` results with `||=`.
+ end
+ RUBY
+
+ expect_correction(<<~RUBY)
+ def foo
+ return @current_user if defined?(@current_user)
 
  @current_user = User.find_by(id: session[:user_id])
+ end
+ RUBY
+ end
+
+ it 'registers registers an offense when the method contains other code' do
+ expect_offense(<<~RUBY)
+ def foo
+ @current_user ||= User.find_by(id: session[:user_id])
+ ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ Avoid memoizing `find_by` results with `||=`.
+ 
+ @current_user.do_something!
+ end
+ RUBY
+
+ expect_correction(<<~RUBY)
+ def foo
+ if defined?(@current_user)
+ @current_user
+ else
+ @current_user = User.find_by(id: session[:user_id])
+ end
+ 
+ @current_user.do_something!
+ end
  RUBY
  end
  end