add Spy#wrap

lib/spy/api.rb

@@ -2,23 +2,22 @@
 
 module Spy
   module API
-    # Initializes a new spy instance for the method
+    # Spies on calls to a method made on an object
     #
-    # With two args:
-    # @param receiver - the receiver of the message you want to spy on
-    # @param msg - the message passed to the receiver that you want to spy on
-    def on(*args)
-      case args.length
-      when 2
-        spied, msg = *args
-        return core.add_spy(spied, spied.method(msg))
-      end
-      raise ArgumentError
+    # @param target - the object you want to spy on
+    # @param msg - the name of the method to spy on
+    def on(target, msg)
+      core.add_spy(target, target.method(msg))
     end
 
-    # TODO docs
-    def on_any_instance(spied, msg)
-      core.add_spy(spied, spied.instance_method(msg))
+    # Spies on calls to a method made on any instance of some class or module
+    #
+    # @param target - class or module to spy on
+    # @param msg - name of the method to spy on
+    # @returns [Spy::Instance]
+    def on_any_instance(target, msg)
+      raise ArgumentError unless target.respond_to?(:instance_method)
+      core.add_spy(target, target.instance_method(msg))
     end
 
     # Stops spying on the method and restores its original functionality
@@ -36,15 +35,16 @@ def on_any_instance(spied, msg)
     def restore(*args)
       case args.length
       when 1
-        return core.remove_all_spies if args.first == :all
+        core.remove_all_spies if args.first == :all
       when 2
-        spied, msg = *args
-        return core.remove_spy(spied, spied.method(msg))
+        target, msg = *args
+        core.remove_spy(target, target.method(msg))
       when 3
-        spied, msg, method_type = *args
-        return core.remove_spy(spied, spied.send(method_type, msg))
+        target, msg, method_type = *args
+        core.remove_spy(target, target.send(method_type, msg))
+      else
+        raise ArgumentError
       end
-      raise ArgumentError
     end
 
     private

lib/spy/instance.rb

@@ -15,6 +15,7 @@ def initialize(spied, original)
       @conditional_filters = []
       @before_callbacks = []
       @after_callbacks = []
+      @around_procs = []
       @call_count = 0
       @call_history = []
       @strategy = Strategy.factory_build(self)
@@ -47,6 +48,13 @@ def when(&block)
         self
       end
 
+      # Expect block to yield. Call the rest of the chain
+      # when it does
+      def wrap(&block)
+        @around_procs << block
+        self
+      end
+
       def before(&block)
         @before_callbacks << block
         self
@@ -82,11 +90,23 @@ def call(context, *args)
 
         if is_active
           @before_callbacks.each {|f| f.call(*args)}
-          @call_count += 1
-          @call_history << MethodCall.new(context, *args)
         end
 
-        result = call_original(context, *args)
+        result = if @around_procs
+                   # Procify the original call
+                   original_proc = Proc.new do
+                     track_call(context, *args) if is_active
+                     call_original(context, *args)
+                   end
+
+                   # Keep wrapping the original proc with each around_proc
+                   @around_procs.reduce(original_proc) do |p, wrapper|
+                     Proc.new { wrapper.call context, *args, &p }
+                   end.call
+                 else
+                   track_call(context, *args) if is_active
+                   call_original(context, *args)
+                 end
 
         if is_active
           @after_callbacks.each {|f| f.call(*args)}
@@ -101,6 +121,11 @@ def call(context, *args)
 
     private
 
+    def track_call(context, *args)
+      @call_count += 1
+      @call_history << MethodCall.new(context, *args)
+    end
+
     def call_original(context, *args)
       if original.is_a?(UnboundMethod)
         original.bind(context).call(*args)

test/spy_test.rb

@@ -144,6 +144,15 @@ def eval_option(opt, *args)
 
   describe 'any_instance' do
     describe 'Spy.on_any_instance' do
+      describe 'an instance' do
+        it 'throws an ArgumentError' do
+          obj = Object.new
+          assert_raises ArgumentError do
+            Spy.on_any_instance(obj, :hello)
+          end
+        end
+      end
+
       # Wrapping
       [
         { name: 'a class and a class-owned method',   to_spy: Proc.new { TestClass },   msg: :class_owned_method },

test/wrap_test.rb

@@ -0,0 +1,63 @@
+require 'test_helper'
+
+class WrapTest < Minitest::Spec
+  class TestClass
+    attr_accessor :string
+
+    def initialize
+      @string = ''
+    end
+
+    def append(char)
+      @string << char
+    end
+  end
+
+  describe 'Spy#wrap' do
+    describe 'followed by the method call' do
+      it 'correctly wraps the call based on the block.call placement' do
+        # yield before
+        spied = TestClass.new
+        spy = Spy.on(spied, :append)
+        spy.wrap do |context, &block|
+          context.string << 'a'
+          block.call
+        end
+
+        spied.append('b')
+
+        assert_equal 'ab', spied.string,
+          'expected wrapping block code to be called before block.call'
+
+        # yield after
+        spied = TestClass.new
+        spy = Spy.on(spied, :append)
+        spy.wrap do |context, &block|
+          block.call
+          context.string << 'a'
+        end
+
+        spied.append('b')
+
+        assert_equal 'ba', spied.string,
+          'expected wrapping block code to be called after block.call'
+      end
+
+      it 'still updates the call count properly even with multiple wraps' do
+        spied = TestClass.new
+        spy = Spy.on(spied, :append)
+        2.times { spy.wrap { |&block| block.call }}
+        spied.append 'a'
+        assert_equal 1, spy.call_count
+      end
+
+      it 'only updates the call count when the actual original call is made' do
+        spied = TestClass.new
+        spy = Spy.on(spied, :append)
+        spy.wrap {}
+        spied.append('b')
+        assert_equal 0, spy.call_count
+      end
+    end
+  end
+end