From f97d47f694f6495b36aba2f1a06510f774bf6cf0 Mon Sep 17 00:00:00 2001
From: Lin Jen-Shin <godfat@godfat.org>
Date: Wed, 11 Apr 2018 00:08:22 +0800
Subject: [PATCH] [mocks] 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 https://github.com/rspec/rspec-mocks/issues/1213

---
This commit was imported from https://github.com/rspec/rspec-mocks/commit/604dc9500b93ff2b805038d4e59146cb5cdf8d2f.
---
 rspec-mocks/Changelog.md                      |  1 +
 .../rspec/mocks/instance_method_stasher.rb    |  7 ++++-
 rspec-mocks/lib/rspec/mocks/method_double.rb  | 18 ++++++++----
 rspec-mocks/spec/rspec/mocks/stub_spec.rb     | 28 +++++++++++++++++++
 4 files changed, 47 insertions(+), 7 deletions(-)

diff --git a/rspec-mocks/Changelog.md b/rspec-mocks/Changelog.md
index 7569c241c..70125b596 100644
--- a/rspec-mocks/Changelog.md
+++ b/rspec-mocks/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)
diff --git a/rspec-mocks/lib/rspec/mocks/instance_method_stasher.rb b/rspec-mocks/lib/rspec/mocks/instance_method_stasher.rb
index ceedb00d7..ce166e8ef 100644
--- a/rspec-mocks/lib/rspec/mocks/instance_method_stasher.rb
+++ b/rspec-mocks/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
diff --git a/rspec-mocks/lib/rspec/mocks/method_double.rb b/rspec-mocks/lib/rspec/mocks/method_double.rb
index a49d91417..5951964f9 100644
--- a/rspec-mocks/lib/rspec/mocks/method_double.rb
+++ b/rspec-mocks/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)
diff --git a/rspec-mocks/spec/rspec/mocks/stub_spec.rb b/rspec-mocks/spec/rspec/mocks/stub_spec.rb
index 5fb6237b5..572b95b03 100644
--- a/rspec-mocks/spec/rspec/mocks/stub_spec.rb
+++ b/rspec-mocks/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