Implement auto-correct for factory calls with block

Maxim Krizhanovski · Maxim Krizhanovski · commit 05d7bb0a8762 · 2019-01-26T20:31:41.000+02:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -3,6 +3,7 @@
 ## Master (Unreleased)
 
 * Add `RSpec/Yield` cop, suggesting using the `and_yield` method when stubbing a method, accepting a block. ([@Darhazer][])
+* Fix `FactoryBot/CreateList` autocorrect crashing when the factory is called with a block=. ([@Darhazer][])
 
 ## 1.31.0 (2019-01-02)
 
diff --git a/lib/rubocop/cop/rspec/factory_bot/create_list.rb b/lib/rubocop/cop/rspec/factory_bot/create_list.rb
@@ -69,9 +69,9 @@ def on_send(node)
 
           def autocorrect(node)
             if style == :create_list
-              autocorrect_n_times_to_create_list(node)
+              CreateListCorrector.new(node)
             else
-              autocorrect_create_list_to_n_times(node)
+              TimesCorrector.new(node)
             end
           end
 
@@ -85,63 +85,127 @@ def contains_only_factory?(node)
             end
           end
 
-          def autocorrect_n_times_to_create_list(node)
-            block = node.parent
-            count = block.receiver.source
-            replacement = factory_call_replacement(block.body, count)
+          # :nodoc
+          class Corrector
+            private
 
-            lambda do |corrector|
-              corrector.replace(block.loc.expression, replacement)
+            def build_options_string(options)
+              options.map(&:source).join(', ')
+            end
+
+            def format_method_call(node, method, arguments)
+              if node.block_type? || node.parenthesized?
+                "#{method}(#{arguments})"
+              else
+                "#{method} #{arguments}"
+              end
+            end
+
+            def format_receiver(receiver)
+              return '' unless receiver
+
+              "#{receiver.source}."
             end
           end
 
-          def autocorrect_create_list_to_n_times(node)
-            replacement = generate_n_times_block(node)
-            lambda do |corrector|
+          # :nodoc
+          class TimesCorrector < Corrector
+            def initialize(node)
+              @node = node
+            end
+
+            def call(corrector)
+              replacement = generate_n_times_block(node)
               corrector.replace(node.loc.expression, replacement)
             end
-          end
 
-          def generate_n_times_block(node)
-            receiver, factory, count, options = *factory_list_call(node)
+            private
 
-            arguments = ":#{factory}"
-            options = build_options_string(options)
-            arguments += ", #{options}" unless options.empty?
+            attr_reader :node
 
-            replacement = format_receiver(receiver)
-            replacement += format_method_call(node, 'create', arguments)
-            "#{count}.times { #{replacement} }"
+            def generate_n_times_block(node)
+              factory, count, *options = node.arguments
+
+              arguments = factory.source
+              options = build_options_string(options)
+              arguments += ", #{options}" unless options.empty?
+
+              replacement = format_receiver(node.receiver)
+              replacement += format_method_call(node, 'create', arguments)
+              "#{count.source}.times { #{replacement} }"
+            end
           end
 
-          def factory_call_replacement(body, count)
-            receiver, factory, options = *factory_call(body)
+          # :nodoc:
+          class CreateListCorrector < Corrector
+            def initialize(node)
+              @node = node.parent
+            end
 
-            arguments = ":#{factory}, #{count}"
-            options = build_options_string(options)
-            arguments += ", #{options}" unless options.empty?
+            def call(corrector)
+              replacement = if node.body.block_type?
+                              call_with_block_replacement(node)
+                            else
+                              call_replacement(node)
+                            end
 
-            replacement = format_receiver(receiver)
-            replacement += format_method_call(body, 'create_list', arguments)
-            replacement
-          end
+              corrector.replace(node.loc.expression, replacement)
+            end
 
-          def build_options_string(options)
-            options.map(&:source).join(', ')
-          end
+            private
 
-          def format_method_call(node, method, arguments)
-            if node.parenthesized?
-              "#{method}(#{arguments})"
-            else
-              "#{method} #{arguments}"
+            attr_reader :node
+
+            def call_with_block_replacement(node)
+              block = node.body
+              arguments = build_arguments(block, node.receiver.source)
+              replacement = format_receiver(block.send_node.receiver)
+              replacement += format_method_call(block, 'create_list', arguments)
+              replacement += format_block(block)
+              replacement
             end
-          end
 
-          def format_receiver(receiver)
-            return '' unless receiver
+            def build_arguments(node, count)
+              factory, *options = *node.send_node.arguments
+
+              arguments = ":#{factory.value}, #{count}"
+              options = build_options_string(options)
+              arguments += ", #{options}" unless options.empty?
+              arguments
+            end
+
+            def call_replacement(node)
+              block = node.body
+              factory, *options = *block.arguments
+
+              arguments = "#{factory.source}, #{node.receiver.source}"
+              options = build_options_string(options)
+              arguments += ", #{options}" unless options.empty?
 
-            "#{receiver.source}."
+              replacement = format_receiver(block.receiver)
+              replacement += format_method_call(block, 'create_list', arguments)
+              replacement
+            end
+
+            def format_block(node)
+              if node.body.begin_type?
+                format_multiline_block(node)
+              else
+                format_singeline_block(node)
+              end
+            end
+
+            def format_multiline_block(node)
+              indent = ' ' * node.body.loc.column
+              indent_end = ' ' * node.parent.loc.column
+              " do #{node.arguments.source}\n" \
+              "#{indent}#{node.body.source}\n" \
+              "#{indent_end}end"
+            end
+
+            def format_singeline_block(node)
+              " { #{node.arguments.source} #{node.body.source} }"
+            end
           end
         end
       end
diff --git a/spec/rubocop/cop/rspec/factory_bot/create_list_spec.rb b/spec/rubocop/cop/rspec/factory_bot/create_list_spec.rb
@@ -88,6 +88,36 @@
     include_examples 'autocorrect',
                      '5.times { FactoryGirl.create :user }',
                      'FactoryGirl.create_list :user, 5'
+
+    bad_code = <<-RUBY
+      3.times do
+        create(:user, :trait) { |user| create :account, user: user }
+      end
+    RUBY
+
+    good_code = <<-RUBY
+      create_list(:user, 3, :trait) { |user| create :account, user: user }
+    RUBY
+
+    include_examples 'autocorrect', bad_code, good_code
+
+    bad_code = <<-RUBY
+      3.times do
+        create(:user, :trait) do |user|
+          create :account, user: user
+          create :profile, user: user
+        end
+      end
+    RUBY
+
+    good_code = <<-RUBY
+      create_list(:user, 3, :trait) do |user|
+          create :account, user: user
+          create :profile, user: user
+      end
+    RUBY
+
+    include_examples 'autocorrect', bad_code, good_code
   end
 
   context 'when EnforcedStyle is :n_times' do
