Teach the cop about travel_back { ... }

lib/rubocop/cop/rspec_rails/timecop.rb

@@ -16,10 +16,10 @@ module RSpecRails
       # rather than relying on allowing time to continue to flow.
       #
       # `Timecop.return` should be replaced with `travel_back`, when used
-      # without a block. `travel_back` does not accept a block, so where
-      # `return` is used with a block, it should be replaced by explicitly
-      # calling `freeze_time` with a block, and passing the `time` to
-      # temporarily return to.
+      # without a block. `travel_back` accepts a block starting with Rails 6.1.
+      # For earlier Rails, where `return` is used with a block, it should
+      # be replaced by explicitly calling `freeze_time` with a block, and
+      # passing the `time` to temporarily return to.
       #
       # `Timecop.travel` should be replaced by `travel` or `travel_to` when
       # passed a `duration` or `time`, respectively. As with `Timecop.scale`,
@@ -81,7 +81,7 @@ module RSpecRails
       #
       #   # good
       #   travel_back
-      #   travel_to(time) { assert true }
+      #   travel_back { assert true }
       #
       #   # bad
       #   Timecop.scale(factor)
@@ -158,7 +158,7 @@ def on_timecop_freeze(node, arguments)
 
         def on_timecop_return(node, arguments)
           message =
-            format(RETURN_MESSAGE, replacement: preferred_return_replacement)
+            format(RETURN_MESSAGE, replacement: 'travel_back')
           add_offense(node, message: message) do |corrector|
             autocorrect_return(corrector, node, arguments)
           end
@@ -180,16 +180,20 @@ def autocorrect_freeze(corrector, node, arguments)
         end
 
         def autocorrect_return(corrector, node, _arguments)
-          return if given_block?(node)
+          return if given_block?(node) && !supports_return_with_block?
 
-          corrector.replace(receiver_and_message_range(node),
-                            preferred_return_replacement)
+          corrector.replace(receiver_and_message_range(node), 'travel_back')
         end
 
         def given_block?(node)
           node.parent&.block_type? && node.parent.send_node == node
         end
 
+        # travel_back { ... } was introduced in Rails 6.1
+        def supports_return_with_block?
+          target_rails_version >= 6.1
+        end
+
         def receiver_and_message_range(node)
           node.source_range.with(end_pos: node.location.selector.end_pos)
         end
@@ -199,12 +203,6 @@ def preferred_freeze_replacement
 
           'freeze_time'
         end
-
-        def preferred_return_replacement
-          return 'travel_back' if target_rails_version < 6.0
-
-          'unfreeze_time'
-        end
       end
     end
   end

spec/rubocop/cop/rspec_rails/timecop_spec.rb

@@ -118,63 +118,48 @@
     end
   end
 
-  shared_examples 'return prefers' do |replacement|
+  shared_examples 'return prefers' do
     context 'when given no block' do
-      it "flags, and corrects to `#{replacement}`" do
+      it 'flags, and corrects to `travel_back`' do
         expect_offense(<<~RUBY)
           Timecop.return
-          ^^^^^^^^^^^^^^ Use `#{replacement}` instead of `Timecop.return`
+          ^^^^^^^^^^^^^^ Use `travel_back` instead of `Timecop.return`
         RUBY
 
         expect_correction(<<~RUBY)
-          #{replacement}
+          travel_back
         RUBY
       end
-
-      context 'when inside a block' do
-        it "flags, and corrects to `#{replacement}`" do
-          expect_offense(<<~RUBY)
-            foo { Timecop.return }
-                  ^^^^^^^^^^^^^^ Use `#{replacement}` instead of `Timecop.return`
-          RUBY
-
-          expect_correction(<<~RUBY)
-            foo { #{replacement} }
-          RUBY
-        end
-      end
     end
+  end
 
-    context 'when given a block' do
-      it 'flags, and does not correct' do
+  describe '.return' do
+    context 'with Rails < 6.1', :rails60 do
+      include_examples 'return prefers'
+
+      it 'flags, but does not correct return with a block' do
         expect_offense(<<~RUBY)
           Timecop.return { assert true }
-          ^^^^^^^^^^^^^^ Use `#{replacement}` instead of `Timecop.return`
+          ^^^^^^^^^^^^^^ Use `travel_back` instead of `Timecop.return`
         RUBY
 
         expect_no_corrections
       end
-
-      context 'when inside a block' do
-        it 'flags, and does not correct' do
-          expect_offense(<<~RUBY)
-            foo { Timecop.return { assert true } }
-                  ^^^^^^^^^^^^^^ Use `#{replacement}` instead of `Timecop.return`
-          RUBY
-
-          expect_no_corrections
-        end
-      end
     end
-  end
 
-  describe '.return' do
-    context 'when Rails < 6.0', :rails52 do
-      include_examples 'return prefers', 'travel_back'
-    end
+    context 'with Rails 6.1+', :rails61 do
+      include_examples 'return prefers'
 
-    context 'with Rails 6.0+', :rails60 do
-      include_examples 'return prefers', 'unfreeze_time'
+      it 'flags, and corrects return with a block' do
+        expect_offense(<<~RUBY)
+          Timecop.return { assert true }
+          ^^^^^^^^^^^^^^ Use `travel_back` instead of `Timecop.return`
+        RUBY
+
+        expect_correction(<<~RUBY)
+          travel_back { assert true }
+        RUBY
+      end
     end
   end