Add TestCase#stub_const (#44294)

* Add TestCase#stub_const

* Note the concurrency issue

* Changelog entry

activesupport/CHANGELOG.md

@@ -1,3 +1,7 @@
+*   Add `ActiveSupport::TestCase#stub_const` to stub a constant for the duration of a yield.
+
+    *DHH*
+
 *   Fix `ActiveSupport::EncryptedConfiguration` to be compatible with Psych 4
 
     *Stephen Sugden*

activesupport/lib/active_support/test_case.rb

@@ -10,6 +10,7 @@ require "active_support/testing/declarative"
 require "active_support/testing/isolation"
 require "active_support/testing/constant_lookup"
 require "active_support/testing/time_helpers"
+require "active_support/testing/constant_stubbing"
 require "active_support/testing/file_fixtures"
 require "active_support/testing/parallelization"
 require "active_support/testing/parallelize_executor"
@@ -126,6 +127,7 @@ module ActiveSupport
     prepend ActiveSupport::Testing::SetupAndTeardown
     include ActiveSupport::Testing::Assertions
     include ActiveSupport::Testing::Deprecation
+    include ActiveSupport::Testing::ConstantStubbing
     include ActiveSupport::Testing::TimeHelpers
     include ActiveSupport::Testing::FileFixtures
     extend ActiveSupport::Testing::Declarative

activesupport/lib/active_support/testing/constant_stubbing.rb

@@ -0,0 +1,30 @@
+module ActiveSupport
+  module Testing
+    module ConstantStubbing
+      # Changes the value of a constant for the duration of a block. Example:
+      #
+      #   # World::List::Import::LARGE_IMPORT_THRESHOLD = 5000
+      #   stub_const(World::List::Import, :LARGE_IMPORT_THRESHOLD, 1) do
+      #     assert_equal 1, World::List::Import::LARGE_IMPORT_THRESHOLD
+      #   end
+      #
+      #   assert_equal 5000, World::List::Import::LARGE_IMPORT_THRESHOLD = 5000
+      #
+      # Using this method rather than forcing `World::List::Import::LARGE_IMPORT_THRESHOLD = 5000` prevents
+      # warnings from being thrown, and ensures that the old value is returned after the test has completed.
+      #
+      # Note: Stubbing a const will stub it across all threads. So if you have concurrent threads
+      # (like separate test suites running in parallel) that all depend on the same constant, it's possible
+      # divergent stubbing will trample on each other.
+      def stub_const(klass, constant, new_value)
+        old_value = klass.const_get(constant)
+        klass.send(:remove_const, constant)
+        klass.const_set(constant, new_value)
+        yield
+      ensure
+        klass.send(:remove_const, constant)
+        klass.const_set(constant, old_value)
+      end
+    end
+  end
+end

activesupport/test/test_case_test.rb

@@ -526,3 +526,29 @@ class TestOrderTest < ActiveSupport::TestCase
     assert_equal :random, Class.new(ActiveSupport::TestCase).test_order
   end
 end
+
+
+class ConstStubbable
+  CONSTANT = 1
+end
+
+class TestConstStubbing < ActiveSupport::TestCase
+  test "stubbing a constant temporarily replaces it with a new value" do
+    stub_const(ConstStubbable, :CONSTANT, 2) do
+      assert_equal 2, ConstStubbable::CONSTANT
+    end
+
+    assert_equal 1, ConstStubbable::CONSTANT
+  end
+
+  test "stubbed constant still reset even if exception is raised" do
+    assert_raises(RuntimeError) do
+      stub_const(ConstStubbable, :CONSTANT, 2) do
+        assert_equal 2, ConstStubbable::CONSTANT
+        raise "Exception"
+      end
+    end
+
+    assert_equal 1, ConstStubbable::CONSTANT
+  end
+end