lib/muack/mock.rb in muack-0.5.0 vs lib/muack/mock.rb in muack-0.5.1

@@ -6,11 +6,11 @@
 module Muack
   class Mock < BasicObject
     attr_reader :object
     def initialize object
       @object = object
-      @__mock_ignore = []
+      @__mock_injected = {}
       [:__mock_defis=, :__mock_disps=].each do |m|
         __send__(m, ::Hash.new{ |h, k| h[k] = [] })
       end
     end
 
@@ -20,11 +20,15 @@
     end
 
     # Public API: Define mocked method
     def method_missing msg, *args, &block
       defi = Definition.new(msg, args, block)
-      __mock_inject_method(defi) if __mock_pure?(defi)
+      if injected = __mock_injected[defi.msg]
+        defi.original_method = injected.original_method
+      else
+        __mock_inject_method(defi)
+      end
       __mock_defis_push(defi)
       Modifier.new(self, defi)
     end
 
     # used for Muack::Modifier#times
@@ -35,15 +39,10 @@
     # used for Muack::Modifier#times
     def __mock_defis_pop defi
       __mock_defis[defi.msg].pop
     end
 
-    # used for Muack::Modifier#times
-    def __mock_ignore defi
-      @__mock_ignore << defi
-    end
-
     # used for mocked object to dispatch mocked method
     def __mock_dispatch msg, actual_args, actual_block
       if defi = __mock_defis[msg].shift
         __mock_disps_push(defi)
         if __mock_check_args(defi.args, actual_args)
@@ -75,33 +74,27 @@
       end
     end
 
     # used for Muack::Session#reset
     def __mock_reset
-      [__mock_defis.values, __mock_disps.values, @__mock_ignore].
-      flatten.compact.each do |defi|
+      __mock_injected.each_value do |defi|
         object.singleton_class.module_eval do
-          methods = instance_methods(false)
-          if methods.include?(defi.msg) # removed mocked method
-            remove_method(defi.msg)     # could be removed by other defi
-          end
-          if methods.include?(defi.original_method) # restore original method
+          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
           end
         end
       end
     end
 
     protected # get warnings for private attributes
-    attr_accessor :__mock_defis, :__mock_disps
+    attr_accessor :__mock_defis, :__mock_disps, :__mock_injected
 
     private
-    def __mock_pure? defi
-      __mock_defis[defi.msg].empty? && __mock_disps[defi.msg].empty?
-    end
-
     def __mock_inject_method defi
+      __mock_injected[defi.msg] = defi
       target = object.singleton_class
       Mock.store_original_method(target, defi)
        __mock_inject_mock_method(target, defi)
     end