[Fix #1062] Correct rake tasks with arguments

changelog/fix_autocorrection_for_rake_environment_with_arguments.md

@@ -0,0 +1 @@
+* [#1062](https://github.com/rubocop/rubocop-rails/issues/1062): Fix autocorrection for `Rails/RakeEnvironment` when rake task accepts arguments. ([@fastjames][])

lib/rubocop/cop/rails/rake_environment.rb

@@ -45,16 +45,24 @@ def on_block(node) # rubocop:disable InternalAffairs/NumblockHandler
             return if with_dependencies?(task_method)
 
             add_offense(task_method) do |corrector|
-              task_name = task_method.arguments[0]
-              task_dependency = correct_task_dependency(task_name)
-
-              corrector.replace(task_name, task_dependency)
+              if with_arguments?(task_method)
+                new_task_dependency = correct_task_arguments_dependency(task_method)
+                corrector.replace(task_arguments(task_method), new_task_dependency)
+              else
+                task_name = task_method.first_argument
+                new_task_dependency = correct_task_dependency(task_name)
+                corrector.replace(task_name, new_task_dependency)
+              end
             end
           end
         end
 
         private
 
+        def correct_task_arguments_dependency(task_method)
+          "#{task_arguments(task_method).source} => :environment"
+        end
+
         def correct_task_dependency(task_name)
           if task_name.sym_type?
             "#{task_name.source.delete(':|\'|"')}: :environment"
@@ -80,6 +88,14 @@ def task_name(node)
           end
         end
 
+        def task_arguments(node)
+          node.arguments[1]
+        end
+
+        def with_arguments?(node)
+          node.arguments.size > 1 && node.arguments[1].array_type?
+        end
+
         def with_dependencies?(node)
           first_arg = node.arguments[0]
           return false unless first_arg

spec/rubocop/cop/rails/rake_environment_spec.rb

@@ -40,6 +40,19 @@
     RUBY
   end
 
+  it 'registers an offense for a task with arguments' do
+    expect_offense(<<~RUBY)
+      task :foo, [:arg] do
+      ^^^^^^^^^^^^^^^^^ Include `:environment` task as a dependency for all Rake tasks.
+      end
+    RUBY
+
+    expect_correction(<<~RUBY)
+      task :foo, [:arg] => :environment do
+      end
+    RUBY
+  end
+
   it 'does not register an offense to task with :environment but it has other dependency before it' do
     expect_no_offenses(<<~RUBY)
       task foo: [:bar, `:environment`] do