Added handling of blocks for mutated instance methods.

lib/chaser.rb

@@ -142,8 +142,9 @@ def record_passing_mutation
   def unmodify_instance_method
     chaser = self
     @mutated = false
-    @klass.send(:define_method, @method_name) do |*args|
-      chaser.old_method.bind(self).call(*args)
+    chaser_proxy_method_name = "__chaser_proxy__#{@method_name}"
+    @klass.send(:define_method, chaser_proxy_method_name) do |block, *args|
+      chaser.old_method.bind(self).call(*args) {|*yielded_values| block.call(*yielded_values)}
     end
   end
 
@@ -155,12 +156,28 @@ def unmodify_class_method
     end
   end
 
+  # Ruby 1.8 doesn't allow define_method to handle blocks.
+  # The blog post http://coderrr.wordpress.com/2008/10/29/using-define_method-with-blocks-in-ruby-18/
+  # show that define_method has problems, and showed how to do workaround_method_code_string
   def modify_instance_method
     chaser = self
     @mutated = true
     @old_method = @klass.instance_method(@method_name)
-    @klass.send(:define_method, @method_name) do |*args|
-      chaser.mutate_value(chaser.old_method.bind(self).call(*args))
+    chaser_proxy_method_name = "__chaser_proxy__#{@method_name}"
+    workaround_method_code_string = <<-EOM
+      def #{@method_name}(*args, &block)
+        #{chaser_proxy_method_name}(block, *args)
+      end
+    EOM
+    @klass.class_eval do
+      eval(workaround_method_code_string)
+    end
+    @klass.send(:define_method, chaser_proxy_method_name) do |block, *args|
+      original_value = chaser.old_method.bind(self).call(*args) do |*yielded_values|
+        mutated_yielded_values = yielded_values.map{|value| chaser.mutate_value(value)}
+        block.call(*mutated_yielded_values)
+      end
+      chaser.mutate_value(original_value)
     end
   end
 

test/fixtures/chased.rb

@@ -10,4 +10,18 @@ def say_hello
   def self.static_method
     "Zap!"
   end
+
+  def block_using_instance_method
+    result = []
+    block_yielding_instance_method do |i|
+      result << i * 2
+    end
+    result
+  end
+
+  def block_yielding_instance_method
+    yield 1
+    yield 2
+    yield 3
+  end
 end

test/test_chaser.rb

@@ -71,6 +71,16 @@ def test_modify_and_unmodify_class_method
     assert_equal "Zap!", Chased.static_method, "class method should be back to normal, but it isn't"
   end
 
+  def test_pass_blocks_on_in_instance_methods
+    @chaser = TestChaser.new("Chased", "block_yielding_instance_method")
+    chased = Chased.new
+    assert_equal [2,4,6], chased.block_using_instance_method, "block yielding instance method has been modified before it should have been"
+    @chaser.modify_method
+    assert_equal [12, 14, 16], chased.block_using_instance_method, "yielded values haven't been modified"
+    @chaser.unmodify_method
+    assert_equal [2,4,6], chased.block_using_instance_method, "block yielding instance method has been modified before it should have been"
+  end
+
 end