stub_const polishing

Author: fxn

activesupport/lib/active_support/testing/constant_stubbing.rb

@@ -10,22 +10,22 @@ module ConstantStubbing
       #     assert_equal 1, World::List::Import::LARGE_IMPORT_THRESHOLD
       #   end
       #
-      #   assert_equal 5000, World::List::Import::LARGE_IMPORT_THRESHOLD = 5000
+      #   assert_equal 5000, World::List::Import::LARGE_IMPORT_THRESHOLD
       #
       # Using this method rather than forcing <tt>World::List::Import::LARGE_IMPORT_THRESHOLD = 5000</tt> 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)
+      def stub_const(mod, constant, new_value)
+        old_value = mod.const_get(constant, false)
+        mod.send(:remove_const, constant)
+        mod.const_set(constant, new_value)
         yield
       ensure
-        klass.send(:remove_const, constant)
-        klass.const_set(constant, old_value)
+        mod.send(:remove_const, constant)
+        mod.const_set(constant, old_value)
       end
     end
   end

activesupport/test/test_case_test.rb

@@ -557,6 +557,9 @@ class ConstStubbable
   CONSTANT = 1
 end
 
+class SubclassOfConstStubbable < ConstStubbable
+end
+
 class TestConstStubbing < ActiveSupport::TestCase
   test "stubbing a constant temporarily replaces it with a new value" do
     stub_const(ConstStubbable, :CONSTANT, 2) do
@@ -576,4 +579,14 @@ class TestConstStubbing < ActiveSupport::TestCase
 
     assert_equal 1, ConstStubbable::CONSTANT
   end
+
+  test "trying to stub a constant that does not exist in the receiver raises NameError" do
+    assert_raises(NameError) do
+      stub_const(ConstStubbable, :NOT_A_CONSTANT, 1) { }
+    end
+
+    assert_raises(NameError) do
+      stub_const(SubclassOfConstStubbable, :CONSTANT, 1) { }
+    end
+  end
 end