lib/rspec/mocks/instance_method_stasher.rb in rspec-mocks-2.99.4 vs lib/rspec/mocks/instance_method_stasher.rb in rspec-mocks-3.0.0.beta1

@@ -1,93 +1,111 @@
 module RSpec
   module Mocks
     # @private
     class InstanceMethodStasher
-      def initialize(klass, method)
-        @klass = klass
+      def initialize(object, method)
+        @object = object
         @method = method
+        @klass = (class << object; self; end)
 
+        @original_method = nil
         @method_is_stashed = false
       end
 
-      # @private
-      def method_is_stashed?
-        @method_is_stashed
-      end
+      attr_reader :original_method
 
-      # @private
-      def stash
-        return if !method_defined_directly_on_klass? || @method_is_stashed
+      if RUBY_VERSION.to_f < 1.9
+        # @private
+        def method_is_stashed?
+          @method_is_stashed
+        end
 
-        @klass.__send__(:alias_method, stashed_method_name, @method)
-        @method_is_stashed = true
-      end
+        # @private
+        def stash
+          return if !method_defined_directly_on_klass? || @method_is_stashed
 
-      private
+          @klass.__send__(:alias_method, stashed_method_name, @method)
+          @method_is_stashed = true
+        end
 
-      # @private
-      def method_defined_directly_on_klass?
-        method_defined_on_klass? && method_owned_by_klass?
-      end
+        # @private
+        def stashed_method_name
+          "obfuscated_by_rspec_mocks__#{@method}"
+        end
+        private :stashed_method_name
 
-      # @private
-      def method_defined_on_klass?(klass = @klass)
-        klass.method_defined?(@method) || klass.private_method_defined?(@method)
-      end
-
-      if ::UnboundMethod.method_defined?(:owner)
         # @private
-        def method_owned_by_klass?
-          owner = @klass.instance_method(@method).owner
+        def restore
+          return unless @method_is_stashed
 
-          # On Ruby 2.0.0+ the owner of a method on a class which has been
-          # `prepend`ed may actually be an instance, e.g.
-          # `#<MyClass:0x007fbb94e3cd10>`, rather than the expected `MyClass`.
-          owner = owner.class unless owner.is_a? Class
-
-          # On 1.8 (and some 1.9s -- e.g. rubinius) aliased methods
-          # can report the wrong owner. Example:
-          # class MyClass
-          #   class << self
-          #     alias alternate_new new
-          #   end
-          # end
-          #
-          # MyClass.owner(:alternate_new) returns `Class` on 1.8,
-          # but we need to consider the owner to be `MyClass` because
-          # it is not actually available on `Class` but is on `MyClass`.
-          # Hence, we verify that the owner actually has the method defined.
-          # If the given owner does not have the method defined, we assume
-          # that the method is actually owned by @klass.
-          owner == @klass || !(method_defined_on_klass?(owner))
+          if @klass.__send__(:method_defined?, @method)
+            @klass.__send__(:undef_method, @method)
+          end
+          @klass.__send__(:alias_method, @method, stashed_method_name)
+          @klass.__send__(:remove_method, stashed_method_name)
+          @method_is_stashed = false
         end
       else
+
         # @private
-        def method_owned_by_klass?
-          # On 1.8.6, which does not support Method#owner, we have no choice but
-          # to assume it's defined on the klass even if it may be defined on
-          # a superclass.
-          true
+        def method_is_stashed?
+          !!@original_method
         end
+
+        # @private
+        def stash
+          return if !method_defined_directly_on_klass?
+          @original_method ||= ::RSpec::Mocks.method_handle_for(@object, @method)
+        end
+
+        # @private
+        def restore
+          return unless @original_method
+
+          if @klass.__send__(:method_defined?, @method)
+            @klass.__send__(:undef_method, @method)
+          end
+
+          @klass.__send__(:define_method, @method, @original_method)
+          @original_method = nil
+        end
       end
 
-      public
+      private
 
       # @private
-      def stashed_method_name
-        "obfuscated_by_rspec_mocks__#{@method}"
+      def method_defined_directly_on_klass?
+        method_defined_on_klass? && method_owned_by_klass?
       end
 
       # @private
-      def restore
-        return unless @method_is_stashed
+      def method_defined_on_klass?(klass = @klass)
+        klass.method_defined?(@method) || klass.private_method_defined?(@method)
+      end
 
-        if @klass.__send__(:method_defined?, @method)
-          @klass.__send__(:undef_method, @method)
-        end
-        @klass.__send__(:alias_method, @method, stashed_method_name)
-        @klass.__send__(:remove_method, stashed_method_name)
-        @method_is_stashed = false
+      def method_owned_by_klass?
+        owner = @klass.instance_method(@method).owner
+
+        # On Ruby 2.0.0+ the owner of a method on a class which has been
+        # `prepend`ed may actually be an instance, e.g.
+        # `#<MyClass:0x007fbb94e3cd10>`, rather than the expected `MyClass`.
+        owner = owner.class unless Module === owner
+
+        # On some 1.9s (e.g. rubinius) aliased methods
+        # can report the wrong owner. Example:
+        # class MyClass
+        #   class << self
+        #     alias alternate_new new
+        #   end
+        # end
+        #
+        # MyClass.owner(:alternate_new) returns `Class` when incorrect,
+        # but we need to consider the owner to be `MyClass` because
+        # it is not actually available on `Class` but is on `MyClass`.
+        # Hence, we verify that the owner actually has the method defined.
+        # If the given owner does not have the method defined, we assume
+        # that the method is actually owned by @klass.
+        owner == @klass || !(method_defined_on_klass?(owner))
       end
     end
   end
 end