Fix stubbing prepended only methods

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