From fdb98d218eadd31f7b513a08697677610921f878 Mon Sep 17 00:00:00 2001
From: David Heinemeier Hansson <david@basecamp.com>
Date: Tue, 1 Feb 2022 12:20:06 +0100
Subject: [PATCH] Add TestCase#stub_const (#44294)

* Add TestCase#stub_const

* Note the concurrency issue

* Changelog entry
---
 activesupport/CHANGELOG.md                    |  4 +++
 activesupport/lib/active_support/test_case.rb |  2 ++
 .../testing/constant_stubbing.rb              | 30 +++++++++++++++++++
 activesupport/test/test_case_test.rb          | 26 ++++++++++++++++
 4 files changed, 62 insertions(+)
 create mode 100644 activesupport/lib/active_support/testing/constant_stubbing.rb

diff --git a/activesupport/CHANGELOG.md b/activesupport/CHANGELOG.md
index 9ba49207c1..907db97d3b 100644
--- a/activesupport/CHANGELOG.md
+++ b/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*
diff --git a/activesupport/lib/active_support/test_case.rb b/activesupport/lib/active_support/test_case.rb
index a2b205c14b..289b152f8e 100644
--- a/activesupport/lib/active_support/test_case.rb
+++ b/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
diff --git a/activesupport/lib/active_support/testing/constant_stubbing.rb b/activesupport/lib/active_support/testing/constant_stubbing.rb
new file mode 100644
index 0000000000..a30a1ea32e
--- /dev/null
+++ b/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
diff --git a/activesupport/test/test_case_test.rb b/activesupport/test/test_case_test.rb
index 68dfe48412..ae5d080940 100644
--- a/activesupport/test/test_case_test.rb
+++ b/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