From c364a0a1dde177ed1a07379fbd49eef36ac40fa4 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] 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
---
 lib/rspec/mocks/instance_method_stasher.rb |  4 +++-
 lib/rspec/mocks/method_double.rb           | 16 ++++++++++------
 spec/rspec/mocks/stub_spec.rb              | 13 +++++++++++++
 3 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/lib/rspec/mocks/instance_method_stasher.rb b/lib/rspec/mocks/instance_method_stasher.rb
index 12edec2fa..f83b3930a 100644
--- a/lib/rspec/mocks/instance_method_stasher.rb
+++ b/lib/rspec/mocks/instance_method_stasher.rb
@@ -139,7 +139,9 @@ def method_owned_by_klass?
         # The owner of M.b is the raw Module object, instead of the expected
         # singleton class of the module
         return true if RUBY_VERSION < '1.9' && owner == @object
-        owner == @klass || !(method_defined_on_klass?(owner))
+        owner == @klass ||
+          owner.singleton_class == @klass || # When `extend self` is used
+          !(method_defined_on_klass?(owner))
       end
     end
   end
diff --git a/lib/rspec/mocks/method_double.rb b/lib/rspec/mocks/method_double.rb
index 4ffdcdf26..0f95aff5c 100644
--- a/lib/rspec/mocks/method_double.rb
+++ b/lib/rspec/mocks/method_double.rb
@@ -83,8 +83,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
@@ -102,10 +105,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
@@ -261,6 +261,10 @@ def definition_target
 
     private
 
+      def method_owner
+        Object.instance_method(:method).bind(object).call(@method_name).owner
+      end
+
       def remove_method_from_definition_target
         definition_target.__send__(:remove_method, @method_name)
       rescue NameError
diff --git a/spec/rspec/mocks/stub_spec.rb b/spec/rspec/mocks/stub_spec.rb
index efa82d37a..6f9f14718 100644
--- a/spec/rspec/mocks/stub_spec.rb
+++ b/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