lib/spy/subroutine.rb in spy-1.0.0 vs lib/spy/subroutine.rb in spy-1.0.1

@@ -50,13 +50,26 @@
 
       if original_method_visibility || !hook_opts[:force]
         @original_method = current_method
       end
 
-      define_method_with = singleton_method ? :define_singleton_method : :define_method
-      base_object.send(define_method_with, method_name, override_method)
+      if original_method && original_method.owner == base_object
+        original_method.owner.send(:remove_method, method_name)
+      end
 
+      if singleton_method
+        if base_object.singleton_class.method_defined?(method_name) || base_object.singleton_class.private_method_defined?(method_name)
+          base_object.singleton_class.alias_method(method_name, method_name)
+        end
+        base_object.define_singleton_method(method_name, override_method)
+      else
+        if base_object.method_defined?(method_name) || base_object.private_method_defined?(method_name)
+          base_object.alias_method(method_name, method_name)
+        end
+        base_object.define_method(method_name, override_method)
+      end
+
       if [:public, :protected, :private].include? hook_opts[:visibility]
         method_owner.send(hook_opts[:visibility], method_name)
       end
 
       Agency.instance.recruit(self)
@@ -67,15 +80,14 @@
     # unhooks method from object
     # @return [self]
     def unhook
       raise NeverHookedError, "'#{method_name}' method has not been hooked" unless hooked?
 
+      method_owner.send(:remove_method, method_name)
       if original_method && method_owner == original_method.owner
-        method_owner.send(:define_method, method_name, original_method)
-        method_owner.send(original_method_visibility, method_name) if original_method_visibility
-      else
-        method_owner.send(:remove_method, method_name)
+        original_method.owner.send(:define_method, method_name, original_method)
+        original_method.owner.send(original_method_visibility, method_name) if original_method_visibility
       end
 
       clear_method!
       Agency.instance.retire(self)
       self
@@ -103,21 +115,21 @@
     #   spy.and_return(true)
     #   spy.and_return { true }
     #   spy.and_return(force: true) { |invalid_arity| true }
     #
     # @return [self]
-    def and_return(value = nil)
+    def and_return(value = nil, &block)
       @do_not_check_plan_arity = false
 
       if block_given?
         if value.is_a?(Hash) && value.has_key?(:force)
           @do_not_check_plan_arity = !!value[:force]
         elsif !value.nil?
           raise ArgumentError, "value and block conflict. Choose one"
         end
 
-        @plan = Proc.new
+        @plan = block
         check_for_too_many_arguments!(@plan)
       else
         @plan = Proc.new { value }
       end
       self
@@ -134,17 +146,33 @@
     end
 
     # tells the spy to call the original method
     # @return [self]
     def and_call_through
-      @plan = Proc.new do |*args, &block|
-        if original_method
-          original_method.call(*args, &block)
-        else
-          base_object.send(:method_missing, method_name, *args, &block)
+      if @base_object.is_a? Class
+        @plan = Proc.new do |object, *args, &block|
+          if original_method
+            if original_method.is_a? UnboundMethod
+              bound_method = original_method.bind(object)
+              bound_method.call(*args, &block)
+            else
+              original_method.call(*args, &block)
+            end
+          else
+            base_object.send(:method_missing, method_name, *args, &block)
+          end
         end
+      else
+        @plan = Proc.new do |*args, &block|
+          if original_method
+            original_method.call(*args, &block)
+          else
+            base_object.send(:method_missing, method_name, *args, &block)
+          end
+        end
       end
+
       self
     end
 
     # @overload and_raise
     # @overload and_raise(ExceptionClass)
@@ -191,24 +219,32 @@
     end
 
     # check if the method was called with the exact arguments
     # @param args Arguments that should have been sent to the method
     # @return [Boolean]
-    def has_been_called_with?(*args)
+    def has_been_called_with?(*args, &block)
       raise NeverHookedError unless @was_hooked
-      match = block_given? ? Proc.new : proc { |call| call.args == args }
+      match = block_given? ? block : proc { |call| call.args == args }
       calls.any?(&match)
     end
 
     # invoke that the method has been called. You really shouldn't use this
     # method.
     def invoke(object, args, block, called_from)
       check_arity!(args.size)
-      result = if @plan
-                 check_for_too_many_arguments!(@plan)
-                 @plan.call(*args, &block)
-               end
+
+      if base_object.is_a? Class
+        result = if @plan
+                   check_for_too_many_arguments!(@plan)
+                   @plan.call(object, *args, &block)
+                 end
+      else
+        result = if @plan
+                   check_for_too_many_arguments!(@plan)
+                   @plan.call(*args, &block)
+                 end
+      end
     ensure
       calls << CallLog.new(object, called_from, args, block, result)
     end
 
     # reset the call log
@@ -271,9 +307,11 @@
 
     def check_for_too_many_arguments!(block)
       return if @do_not_check_plan_arity || arity_range.nil?
       min_arity = block.arity
       min_arity = min_arity.abs - 1 if min_arity < 0
+
+      min_arity -=1 if base_object.is_a? Class # Instance-method procs take an extra param for receiving object
 
       if min_arity > arity_range.max
         raise ArgumentError.new("block requires #{min_arity} arguments while original_method require a maximum of #{arity_range.max}")
       end
     end