Fix stubbing prepended only methods

Previously, we're assuming the method must be defined in the
singleton class. However this is not always true. Whenever the
method was only defined in the prepended module, then it's not
defined in the singleton class. We need to find the owner of
the method instead, which is the prepended module.

Closes #1213

Author: godfat

Changelog.md

@@ -21,6 +21,7 @@ Bug Fixes:
 
 * Support keyword argument semantics when constraining argument expectations using
   `with` on Ruby 3.0+ (Yusuke Endoh, #1394)
+* Fix stubbing of prepended-only methods. (Lin Jen-Shin, #1218)
 
 ### 3.10.2 / 2021-01-27
 [Full Changelog](http://github.com/rspec/rspec-mocks/compare/v3.10.1...v3.10.2)

lib/rspec/mocks/instance_method_stasher.rb

@@ -62,7 +62,12 @@ def method_owned_by_klass?
         # `#<MyClass:0x007fbb94e3cd10>`, rather than the expected `MyClass`.
         owner = owner.class unless Module === owner
 
-        owner == @klass || !(method_defined_on_klass?(owner))
+        owner == @klass ||
+          # When `extend self` is used, and not under `allow_any_instance_of`
+          # nor `expect_any_instance_of`.
+          (owner.singleton_class == @klass &&
+            !Mocks.space.any_instance_recorder_for(owner, true)) ||
+          !(method_defined_on_klass?(owner))
       end
     end
   end

lib/rspec/mocks/method_double.rb

@@ -85,8 +85,11 @@ def restore_original_method
         return unless @method_is_proxied
 
         remove_method_from_definition_target
-        @method_stasher.restore if @method_stasher.method_is_stashed?
-        restore_original_visibility
+
+        if @method_stasher.method_is_stashed?
+          @method_stasher.restore
+          restore_original_visibility
+        end
 
         @method_is_proxied = false
       end
@@ -104,10 +107,7 @@ def show_frozen_warning
 
       # @private
       def restore_original_visibility
-        return unless @original_visibility &&
-          MethodReference.method_defined_at_any_visibility?(object_singleton_class, @method_name)
-
-        object_singleton_class.__send__(@original_visibility, method_name)
+        method_owner.__send__(@original_visibility, @method_name)
       end
 
       # @private
@@ -249,6 +249,12 @@ def new_rspec_prepended_module
         end
       end
 
+      def method_owner
+        @method_owner ||=
+          # We do this because object.method might be overridden.
+          ::RSpec::Support.method_handle_for(object, @method_name).owner
+      end
+
       def remove_method_from_definition_target
         # In Ruby 2.4 and earlier, `remove_method` is private
         definition_target.__send__(:remove_method, @method_name)

spec/rspec/mocks/stub_spec.rb

@@ -138,6 +138,10 @@ module ToBePrepended
           def value
             "#{super}_prepended".to_sym
           end
+
+          def value_without_super
+            :prepended
+          end
         end
 
         it "handles stubbing prepended methods" do
@@ -165,6 +169,15 @@ def object.value; :original; end
           expect(object.value).to eq :stubbed
         end
 
+        it "handles stubbing prepending methods that were only defined on the prepended module" do
+          object = Object.new
+          object.singleton_class.send(:prepend, ToBePrepended)
+
+          expect(object.value_without_super).to eq :prepended
+          allow(object).to receive(:value_without_super) { :stubbed }
+          expect(object.value_without_super).to eq :stubbed
+        end
+
         it 'does not unnecessarily prepend a module when the prepended module does not override the stubbed method' do
           object = Object.new
           def object.value; :original; end
@@ -350,6 +363,21 @@ class << self; public :hello; end;
           expect(mod.hello).to eq(:hello)
         end
 
+        it "correctly restores from allow_any_instance_of for self extend" do
+          mod = Module.new {
+            extend self
+            def hello; :hello; end
+          }
+
+          allow_any_instance_of(mod).to receive(:hello) { :stub }
+
+          expect(mod.hello).to eq(:stub)
+
+          reset_all
+
+          expect(mod.hello).to eq(:hello)
+        end
+
         it "correctly handles stubbing inherited mixed in class methods" do
           mod = Module.new do
             def method_a