Exclude stubbed classes from subclasses after teardown

lib/rspec/mocks.rb

@@ -113,21 +113,22 @@ class << self
 
     # To speed up boot time a bit, delay loading optional or rarely
     # used features until their first use.
-    autoload :AnyInstance,      "rspec/mocks/any_instance"
-    autoload :ExpectChain,      "rspec/mocks/message_chain"
-    autoload :StubChain,        "rspec/mocks/message_chain"
+    autoload :AnyInstance, "rspec/mocks/any_instance"
+    autoload :ExpectChain, "rspec/mocks/message_chain"
+    autoload :StubChain, "rspec/mocks/message_chain"
     autoload :MarshalExtension, "rspec/mocks/marshal_extension"
+    autoload :ExcludeStubbedClassesFromSubclasses, "rspec/mocks/exclude_stubbed_classes_from_subclasses"
 
     # Namespace for mock-related matchers.
     module Matchers
       # @private
       # just a "tag" for rspec-mock matchers detection
       module Matcher; end
 
-      autoload :HaveReceived,        "rspec/mocks/matchers/have_received"
-      autoload :Receive,             "rspec/mocks/matchers/receive"
+      autoload :HaveReceived, "rspec/mocks/matchers/have_received"
+      autoload :Receive, "rspec/mocks/matchers/receive"
       autoload :ReceiveMessageChain, "rspec/mocks/matchers/receive_message_chain"
-      autoload :ReceiveMessages,     "rspec/mocks/matchers/receive_messages"
+      autoload :ReceiveMessages, "rspec/mocks/matchers/receive_messages"
     end
   end
 end

lib/rspec/mocks/configuration.rb

@@ -10,6 +10,7 @@ def initialize
         @verify_partial_doubles = false
         @temporarily_suppress_partial_double_verification = false
         @color = false
+        @exclude_stubbed_classes_from_subclasses = false
       end
 
       # Sets whether RSpec will warn, ignore, or fail a test when
@@ -178,6 +179,22 @@ def color?
         end
       end
 
+      def exclude_stubbed_classes_from_subclasses?
+        @exclude_stubbed_classes_from_subclasses
+      end
+
+      # When this is set to true, stubbed classes are excluded from the list of
+      # subclasses of the parent class after each spec.
+      def exclude_stubbed_classes_from_subclasses=(val)
+        @exclude_stubbed_classes_from_subclasses = !!val
+
+        if val
+          RSpec::Mocks::ExcludeStubbedClassesFromSubclasses.enable!
+        else
+          RSpec::Mocks::ExcludeStubbedClassesFromSubclasses.disable!
+        end
+      end
+
       # Monkey-patch `Marshal.dump` to enable dumping of mocked or stubbed
       # objects. By default this will not work since RSpec mocks works by
       # adding singleton methods that cannot be serialized. This patch removes

lib/rspec/mocks/exclude_stubbed_classes_from_subclasses.rb

@@ -0,0 +1,56 @@
+module RSpec
+  module Mocks
+    # Support for `exclude_stubbed_classes_from_subclasses` configuration.
+    #
+    # @private
+    class ExcludeStubbedClassesFromSubclasses
+      class << self
+        def enable!
+          return unless RUBY_VERSION >= "3.1"
+          return if Class.respond_to?(:subclasses_with_rspec_mocks)
+
+          Class.class_eval do
+            def subclasses_with_rspec_mocks
+              subclasses_without_rspec_mocks - RSpec::Mocks::ExcludeStubbedClassesFromSubclasses.excluded_subclasses
+            end
+
+            alias subclasses_without_rspec_mocks subclasses
+            alias subclasses subclasses_with_rspec_mocks
+          end
+        end
+
+        def disable!
+          @excluded_subclasses = []
+
+          if Class.respond_to?(:subclasses_with_rspec_mocks)
+            Class.class_eval do
+              undef subclasses_with_rspec_mocks
+              alias subclasses subclasses_without_rspec_mocks
+              undef subclasses_without_rspec_mocks
+            end
+          end
+        end
+
+        def excluded_subclasses
+          require 'weakref' unless defined?(::WeakRef)
+
+          @excluded_subclasses ||= []
+          @excluded_subclasses.select(&:weakref_alive?).map do |ref|
+            begin # rubocop:disable Style/RedundantBegin
+              ref.__getobj__
+            rescue ::WeakRef::RefError
+              nil
+            end
+          end.compact
+        end
+
+        def exclude_subclass(constant)
+          require 'weakref' unless defined?(::WeakRef)
+
+          @excluded_subclasses ||= []
+          @excluded_subclasses << ::WeakRef.new(constant)
+        end
+      end
+    end
+  end
+end

lib/rspec/mocks/mutate_const.rb

@@ -221,6 +221,7 @@ def to_constant
         end
 
         def reset
+          RSpec::Mocks::ExcludeStubbedClassesFromSubclasses.exclude_subclass(@mutated_value) if RSpec::Mocks.configuration.exclude_stubbed_classes_from_subclasses?
           @constants_to_transfer.each do |const|
             @mutated_value.__send__(:remove_const, const)
           end
@@ -297,6 +298,7 @@ def to_constant
         end
 
         def reset
+          RSpec::Mocks::ExcludeStubbedClassesFromSubclasses.exclude_subclass(@mutated_value) if RSpec::Mocks.configuration.exclude_stubbed_classes_from_subclasses?
           @parent.__send__(:remove_const, @const_name)
         end
 

spec/rspec/mocks/exclude_stubbed_classes_from_subclasses_spec.rb

@@ -0,0 +1,97 @@
+if RUBY_VERSION >= '3.1'
+  class TestClass
+  end
+
+  module RSpec
+    module Mocks
+      RSpec.describe ExcludeStubbedClassesFromSubclasses do
+        after do
+          described_class.disable!
+        end
+
+        describe '.enable!' do
+          it 'extends Class with methods' do
+            expect {
+              described_class.enable!
+            }.to change { Class.respond_to?(:subclasses_with_rspec_mocks) }.from(false).to(true)
+          end
+
+          it 'does not extends Class when it has been enabled' do
+            allow(Class).to receive(:class_eval).and_call_original
+
+            described_class.enable!
+            described_class.enable!
+
+            expect(Class).to have_received(:class_eval).once
+          end
+
+          it 'excludes stubbed classes from subclasses' do
+            described_class.enable!
+
+            orignal_subclasses = TestClass.subclasses
+
+            subclass = Class.new(TestClass)
+            described_class.exclude_subclass(subclass)
+
+            expect(TestClass.subclasses).to an_array_matching(orignal_subclasses)
+          end
+        end
+
+        describe '.disable!' do
+          it 'does nothing when it has not been enabled' do
+            expect { described_class.disable! }.not_to raise_error
+          end
+
+          it 'removes methods from class when it has been enabled' do
+            described_class.enable!
+            expect {
+              described_class.disable!
+            }.to change { Class.respond_to?(:subclasses_with_rspec_mocks) }.from(true).to(false)
+          end
+
+          it 'does not exclude stubbed classes from subclasses' do
+            described_class.enable!
+            described_class.disable!
+
+            orignal_subclasses = TestClass.subclasses
+
+            subclass = Class.new(TestClass)
+            described_class.exclude_subclass(subclass)
+
+            expect(TestClass.subclasses).to an_array_matching(orignal_subclasses + [subclass])
+          end
+        end
+
+        describe '.excluded_subclasses' do
+          it 'returns excluded subclasses' do
+            subclass = Class.new
+            described_class.exclude_subclass(subclass)
+
+            expect(described_class.excluded_subclasses).to an_array_matching([subclass])
+          end
+
+          it 'does not return excluded subclasses that have been garbage collected' do
+            subclass = Class.new
+            described_class.exclude_subclass(subclass)
+
+            subclass = nil
+
+            GC.start
+
+            expect(described_class.excluded_subclasses).to eq([])
+          end
+
+          it 'does not return excluded subclasses that raises a ::WeakRef::RefError' do
+            require 'weakref'
+            subclass = double(:weakref_alive? => true)
+            described_class.instance_variable_set(:@excluded_subclasses, [subclass])
+
+            allow(subclass).to receive(:__getobj__).and_raise(::WeakRef::RefError)
+
+            expect(described_class.excluded_subclasses).to eq([])
+          end
+        end
+      end
+    end
+  end
+end

spec/rspec/mocks/mutate_const_spec.rb

@@ -145,6 +145,20 @@ def change_const_value_to(value)
         it 'returns nil' do
           expect(hide_const(const_name)).to be_nil
         end
+
+        if RUBY_VERSION >= '3.1'
+          describe 'with global exclude_stubbed_classes_from_subclasses option set' do
+            include_context "with isolated configuration"
+            include_context "with stubbed classes excluded from subclasses"
+
+            it 'gives the same subclasses after rspec clears its mocks' do
+              original_subclasses = TestClass.subclasses
+              stub_const(const_name, Class.new(TestClass))
+              reset_rspec_mocks
+              expect(TestClass.subclasses).to an_array_matching(original_subclasses)
+            end
+          end
+        end
       end
 
       describe "#hide_const" do
@@ -351,6 +365,20 @@ def change_const_value_to(value)
               end
             end
           end
+
+          if RUBY_VERSION >= '3.1'
+            describe 'with global exclude_stubbed_classes_from_subclasses option set' do
+              include_context "with isolated configuration"
+              include_context "with stubbed classes excluded from subclasses"
+
+              it 'gives the same subclasses after rspec clears its mocks' do
+                original_subclasses = TestClass::Nested.subclasses
+                stub_const('TestClass', Class.new(TestClass::Nested))
+                reset_rspec_mocks
+                expect(TestClass::Nested.subclasses).to an_array_matching(original_subclasses)
+              end
+            end
+          end
         end
 
         context 'for a loaded nested constant' do

spec/spec_helper.rb

@@ -174,3 +174,13 @@ def self.fake_matcher_description
     RSpec::Mocks.configuration.syntax = orig_syntax
   end
 end
+
+RSpec.shared_context "with stubbed classes excluded from subclasses" do
+  before do
+    RSpec::Mocks.configuration.exclude_stubbed_classes_from_subclasses = true
+  end
+
+  after do
+    RSpec::Mocks.configuration.exclude_stubbed_classes_from_subclasses = false
+  end
+end