From 1d40e37a44c6313b89b6d1060bcee3dfe29ccd27 Mon Sep 17 00:00:00 2001
From: r7kamura <r7kamura@gmail.com>
Date: Sat, 27 Jan 2024 10:30:51 +0900
Subject: [PATCH] Fix inline `in_batches` false-positive on
 `Migration/BatchInBatches`

---
 lib/rubocop/cop/migration/batch_in_batches.rb | 29 ++++++++++++++++---
 .../cop/migration/batch_in_batches_spec.rb    | 16 ++++++++--
 2 files changed, 39 insertions(+), 6 deletions(-)

diff --git a/lib/rubocop/cop/migration/batch_in_batches.rb b/lib/rubocop/cop/migration/batch_in_batches.rb
index 28c979d..42fc001 100644
--- a/lib/rubocop/cop/migration/batch_in_batches.rb
+++ b/lib/rubocop/cop/migration/batch_in_batches.rb
@@ -26,6 +26,15 @@ module Migration
       #     disable_ddl_transaction!
       #
       #     def up
+      #       User.in_batches.update_all(some_column: 'some value')
+      #     end
+      #   end
+      #
+      #   # good
+      #   class BackfillSomeColumnToUsers < ActiveRecord::Migration[7.0]
+      #     disable_ddl_transaction!
+      #
+      #     def up
       #       User.in_batches do |relation|
       #         relation.update_all(some_column: 'some value')
       #       end
@@ -75,19 +84,31 @@ def autocorrect(
           )
         end
 
-        # @param node [RuboCop::AST::Node]
+        # @param node [RuboCop::AST::SendNode]
         # @return [Boolean]
-        def within_in_batches?(node)
+        def in_batches?(node)
+          in_block_batches?(node) || in_inline_batches?(node)
+        end
+
+        # @param node [RuboCop::AST::SendNode]
+        # @return [Boolean]
+        def in_block_batches?(node)
           node.each_ancestor(:block).any? do |ancestor|
             ancestor.method?(:in_batches)
           end
         end
 
+        # @param node [RuboCop::AST::SendNode]
+        # @return [Boolean]
+        def in_inline_batches?(node)
+          node.receiver.is_a?(::RuboCop::AST::SendNode) &&
+            node.receiver.method?(:in_batches)
+        end
+
         # @param node [RuboCop::AST::SendNode]
         # @return [Boolean]
         def wrong?(node)
-          batch_processing?(node) &&
-            !within_in_batches?(node)
+          batch_processing?(node) && !in_batches?(node)
         end
       end
     end
diff --git a/spec/rubocop/cop/migration/batch_in_batches_spec.rb b/spec/rubocop/cop/migration/batch_in_batches_spec.rb
index 4e28851..916ac0f 100644
--- a/spec/rubocop/cop/migration/batch_in_batches_spec.rb
+++ b/spec/rubocop/cop/migration/batch_in_batches_spec.rb
@@ -1,7 +1,7 @@
 # frozen_string_literal: true
 
 RSpec.describe RuboCop::Cop::Migration::BatchInBatches, :config do
-  context 'when `update_all` is used within `in_batches`' do
+  context 'when `update_all` is used with block `in_batches`' do
     it 'does not register an offense' do
       expect_no_offenses(<<~RUBY)
         class BackfillUsersSomeColumn < ActiveRecord::Migration[7.0]
@@ -15,7 +15,19 @@ def change
     end
   end
 
-  context 'when `update_all` is used not within `in_batches`' do
+  context 'when `update_all` is used with inline `in_batches`' do
+    it 'does not register an offense' do
+      expect_no_offenses(<<~RUBY)
+        class BackfillUsersSomeColumn < ActiveRecord::Migration[7.0]
+          def change
+            User.in_batches.update_all(some_column: 'some value')
+          end
+        end
+      RUBY
+    end
+  end
+
+  context 'when `update_all` is used not with block `in_batches`' do
     it 'registers an offense' do
       expect_offense(<<~RUBY)
         class BackfillUsersSomeColumn < ActiveRecord::Migration[7.0]