lib/muack/mock.rb in muack-1.5.1 vs lib/muack/mock.rb in muack-1.6.0

@@ -110,59 +110,105 @@
     attr_accessor :__mock_defis, :__mock_disps, :__mock_injected
 
     private
     def __mock_inject_method defi
       __mock_injected[defi.msg] = defi
-      # a) ancestors.first is the first module in the method chain.
-      #    it's just the singleton_class when nothing was prepended,
-      #    otherwise the last prepended module.
-      # b) would be the class in AnyInstanceOf.
-      target = object.singleton_class.ancestors.first
+      target = Mock.prepare_target(object.singleton_class, defi.msg)
+      defi.target = target
       Mock.store_original_method(target, defi)
       __mock_inject_mock_method(target, defi)
     end
 
     def __mock_reset_method defi
-      object.singleton_class.ancestors.first.module_eval do
+      defi.target.module_eval do
         remove_method(defi.msg)
         # restore original method
-        if public_instance_methods(false).include?(defi.original_method) ||
-           protected_instance_methods(false).include?(defi.original_method) ||
-           private_instance_methods(false).include?(defi.original_method)
+        if defi.original_method &&
+           Mock.direct_method_defined?(self, defi.original_method)
           alias_method(defi.msg, defi.original_method)
           __send__(defi.visibility, defi.msg)
           remove_method(defi.original_method)
         end
       end
     end
 
-    def self.store_original_method klass, defi
-      visibility = if klass.public_instance_methods(false).include?(defi.msg)
-        :public
-      elsif klass.protected_instance_methods(false).include?(defi.msg)
-        :protected
-      elsif klass.private_instance_methods(false).include?(defi.msg)
-        :private
+    if ::Class.instance_method(:method_defined?).arity == 1 # Ruby 2.5-
+      def self.direct_method_defined? mod, msg
+        mod.public_instance_methods(false).include?(msg) ||
+          mod.protected_instance_methods(false).include?(msg) ||
+          mod.private_instance_methods(false).include?(msg)
       end
 
-      if visibility # store original method
-        original_method = find_new_name(klass, defi.msg)
-        klass.__send__(:alias_method, original_method, defi.msg)
+      def self.method_visibility mod, msg
+        if mod.public_instance_methods(false).include?(msg)
+          :public
+        elsif mod.protected_instance_methods(false).include?(msg)
+          :protected
+        elsif mod.private_instance_methods(false).include?(msg)
+          :private
+        end
+      end
+    else # Ruby 2.6+
+      def self.direct_method_defined? mod, msg
+        mod.method_defined?(msg, false) || # this doesn't cover private method
+          mod.private_method_defined?(msg, false)
+      end
+
+      def self.method_visibility mod, msg
+        if mod.public_method_defined?(msg, false)
+          :public
+        elsif mod.protected_method_defined?(msg, false)
+          :protected
+        elsif mod.private_method_defined?(msg, false)
+          :private
+        end
+      end
+    end
+
+    def self.prepare_target singleton_class, msg
+      if singleton_class == singleton_class.ancestors.first # no prepended mod
+        singleton_class
+      else # check if we need to prepend an internal module to override
+        index = singleton_class.ancestors.index(singleton_class)
+        prepended_modules = singleton_class.ancestors[0...index]
+
+        if prepended_modules.find{ |m| direct_method_defined?(m, msg) }
+          # prepend an internal module to override the prepended module(s)
+          name = :MuackPrepended
+          if singleton_class.const_defined?(name)
+            singleton_class.const_get(name)
+          else
+            prepended = ::Module.new
+            singleton_class.const_set(name, prepended)
+            singleton_class.private_constant(name)
+            singleton_class.prepend(prepended)
+            prepended
+          end
+        else # we don't need to override so singleton class is enough
+          singleton_class
+        end
+      end
+    end
+
+    def self.store_original_method mod, defi
+      if visibility = method_visibility(mod, defi.msg)
+        original_method = find_new_name(mod, defi.msg)
+        mod.__send__(:alias_method, original_method, defi.msg)
         defi.original_method = original_method
         defi.visibility = visibility
       else
         defi.visibility = :public
       end
     end
 
-    def self.find_new_name klass, message, level=0
+    def self.find_new_name mod, message, level=0
       if level >= (::ENV['MUACK_RECURSION_LEVEL'] || 9).to_i
         raise CannotFindInjectionName.new(level+1, message)
       end
 
       new_name = "__muack_#{name}_#{level}_#{message}".to_sym
-      if klass.instance_methods(false).include?(new_name)
-        find_new_name(klass, message, level+1)
+      if direct_method_defined?(mod, new_name)
+        find_new_name(mod, message, level+1)
       else
         new_name
       end
     end