Accept block in travel_back time helper

activesupport/CHANGELOG.md

@@ -1,3 +1,7 @@
+*   Add block support to `ActiveSupport::Testing::TimeHelpers#travel_back`.
+
+    *Tim Masliuchenko*
+
 *   Update `ActiveSupport::Messages::Metadata#fresh?` to work for cookies with expiry set when
     `ActiveSupport.parse_json_times = true`.
 

activesupport/lib/active_support/testing/time_helpers.rb

@@ -39,6 +39,10 @@ module ActiveSupport
         @stubs[object.object_id][method_name]
       end
 
+      def stubbed?
+        !@stubs.empty?
+      end
+
       private
         def unstub_object(stub)
           singleton_class = stub.object.singleton_class
@@ -160,12 +164,32 @@ module ActiveSupport
       # +travel+, +travel_to+, and +freeze_time+.
       #
       #   Time.current # => Sat, 09 Nov 2013 15:34:49 EST -05:00
+      #
       #   travel_to Time.zone.local(2004, 11, 24, 01, 04, 44)
       #   Time.current # => Wed, 24 Nov 2004 01:04:44 EST -05:00
+      #
       #   travel_back
       #   Time.current # => Sat, 09 Nov 2013 15:34:49 EST -05:00
+      #
+      # This method also accepts a block, which brings the stubs back at the end of the block:
+      #
+      #   Time.current # => Sat, 09 Nov 2013 15:34:49 EST -05:00
+      #
+      #   travel_to Time.zone.local(2004, 11, 24, 01, 04, 44)
+      #   Time.current # => Wed, 24 Nov 2004 01:04:44 EST -05:00
+      #
+      #   travel_back do
+      #     Time.current # => Sat, 09 Nov 2013 15:34:49 EST -05:00
+      #   end
+      #
+      #   Time.current # => Wed, 24 Nov 2004 01:04:44 EST -05:00
       def travel_back
+        stubbed_time = Time.current if block_given? && simple_stubs.stubbed?
+
         simple_stubs.unstub_all!
+        yield if block_given?
+      ensure
+        travel_to stubbed_time if stubbed_time
       end
       alias_method :unfreeze_time, :travel_back
 

activesupport/test/time_travel_test.rb

@@ -102,6 +102,29 @@ class TimeTravelTest < ActiveSupport::TestCase
     end
   end
 
+  def test_time_helper_travel_back_with_block
+    Time.stub(:now, Time.now) do
+      expected_time = Time.new(2004, 11, 24, 01, 04, 44)
+
+      travel_to expected_time
+      assert_equal expected_time, Time.now
+      assert_equal Date.new(2004, 11, 24), Date.today
+      assert_equal expected_time.to_datetime, DateTime.now
+
+      travel_back do
+        assert_not_equal expected_time, Time.now
+        assert_not_equal Date.new(2004, 11, 24), Date.today
+        assert_not_equal expected_time.to_datetime, DateTime.now
+      end
+
+      assert_equal expected_time, Time.now
+      assert_equal Date.new(2004, 11, 24), Date.today
+      assert_equal expected_time.to_datetime, DateTime.now
+    ensure
+      travel_back
+    end
+  end
+
   def test_time_helper_travel_to_with_nested_calls_with_blocks
     Time.stub(:now, Time.now) do
       outer_expected_time = Time.new(2004, 11, 24, 01, 04, 44)