lib/muack/mock.rb in muack-0.7.3 vs lib/muack/mock.rb in muack-1.0.0

@@ -1,12 +1,15 @@
 
 require 'muack/definition'
 require 'muack/modifier'
 require 'muack/failure'
+require 'muack/block'
 require 'muack/error'
 
 module Muack
+  EmptyBlock = proc{}
+
   class Mock < BasicObject
     attr_reader :object
     def initialize object
       @object = object
       @__mock_injected = {}
@@ -68,25 +71,33 @@
         end
       end
     end
 
     # used for mocked object to dispatch mocked method
-    def __mock_dispatch_call disp, actual_args, actual_block
-      if disp.proxy
-        value = yield # need the original context for proxy or AnyInstanceOf
-        if disp.block
-          disp.block.call(value)
-        else
-          value
-        end
-      elsif block = disp.block
-        arity = block.arity
-        if arity < 0
-          block.call(*actual_args             , &actual_block)
-        else
-          block.call(*actual_args.first(arity), &actual_block)
-        end
+    def __mock_dispatch_call context, disp, actual_args, actual_block, &_yield
+      args = if disp.peek_args
+               __mock_block_call(context, disp.peek_args,
+                                 actual_args, actual_block, true)
+             else
+               actual_args
+             end
+
+      ret = if disp.returns
+              __mock_block_call(context, disp.returns,
+                                args, actual_block, true)
+            elsif disp.original_method # proxies for singleton methods
+              context.__send__(disp.original_method, *args, &actual_block)
+            else # proxies for instance methods
+              # need the original context for calling `super`
+              # ruby: can't pass a block to yield, so we name it _yield
+              _yield.call(*args, &actual_block)
+            end
+
+      if disp.peek_return
+        __mock_block_call(context, disp.peek_return, ret, EmptyBlock, false)
+      else
+        ret
       end
     end
 
     # used for Muack::Session#verify
     def __mock_verify
@@ -103,13 +114,14 @@
     def __mock_reset
       __mock_injected.each_value do |defi|
         object.singleton_class.module_eval do
           remove_method(defi.msg)
           # restore original method
-          if instance_methods(false).include?(defi.original_method)
-            alias_method defi.msg, defi.original_method
-            remove_method defi.original_method
+          if instance_methods(false).include?(defi.original_method) ||
+             private_instance_methods(false).include?(defi.original_method)
+            alias_method(defi.msg, defi.original_method)
+            remove_method(defi.original_method)
           end
         end
       end
     end
 
@@ -118,46 +130,64 @@
 
     private
     def __mock_inject_method defi
       __mock_injected[defi.msg] = defi
       target = object.singleton_class # would be the class in AnyInstanceOf
-      Mock.store_original_method(target, defi)
-       __mock_inject_mock_method(target, defi)
+      privilege = Mock.store_original_method(target, defi)
+      __mock_inject_mock_method(target, defi, privilege)
     end
 
     def self.store_original_method klass, defi
-      return unless klass.instance_methods(false).include?(defi.msg)
+      privilege = if klass.instance_methods(false).include?(defi.msg)
+        :public # TODO: forget about protected methods?
+      elsif klass.private_instance_methods(false).include?(defi.msg)
+        :private
+      end
+
+      return :public unless privilege
       # store original method
       original_method = find_new_name(klass, defi.msg)
       klass.__send__(:alias_method, original_method, defi.msg)
       defi.original_method = original_method
+      privilege
     end
 
     def self.find_new_name klass, message, level=0
       if level >= (::ENV['MUACK_RECURSION_LEVEL'] || 9).to_i
         raise CannotFindInjectionName.new(level+1, message)
       end
 
-      new_name = "__muack_mock_#{level}_#{message}".to_sym
+      new_name = "__muack_#{name}_#{level}_#{message}".to_sym
       if klass.instance_methods(false).include?(new_name)
         find_new_name(klass, message, level+1)
       else
         new_name
       end
     end
 
-    def __mock_inject_mock_method target, defi
+    def __mock_inject_mock_method target, defi, privilege=:public
       mock = self # remember the context
-      target.__send__(:define_method, defi.msg){|*actual_args, &actual_block|
+      target.__send__(:define_method, defi.msg){ |*actual_args, &actual_block|
         disp = mock.__mock_dispatch(defi.msg, actual_args)
-        mock.__mock_dispatch_call(disp, actual_args, actual_block){
-          if disp.original_method
-            __send__(disp.original_method, *actual_args, &actual_block)
-          else
-            super(*actual_args, &actual_block)
-          end
-        }
+        mock.__mock_dispatch_call(self, disp, actual_args,
+                                              actual_block) do |args, &block|
+          super(*args, &block)
+        end
       }
+      target.__send__(privilege, defi.msg)
+    end
+
+    # used for __mock_dispatch_call
+    def __mock_block_call context, block, actual_args, actual_block, splat
+      return unless block
+      # for AnyInstanceOf, we don't have actually context at the time
+      # we're defining it, so we update it here
+      block.context = context if block.kind_of?(Block)
+      if splat
+        block.call(*actual_args, &actual_block)
+      else # peek_return doesn't need splat
+        block.call(actual_args, &actual_block)
+      end
     end
 
     def __mock_check_args expected_args, actual_args
       if expected_args == [WithAnyArgs]
         true