lib/rspec/mocks/method_double.rb in rspec-mocks-3.0.0.beta1 vs lib/rspec/mocks/method_double.rb in rspec-mocks-3.0.0.beta2

@@ -33,19 +33,11 @@
 
       alias_method :save_original_method!, :original_method
 
       # @private
       def visibility
-        if TestDouble === @object
-          'public'
-        elsif object_singleton_class.private_method_defined?(@method_name)
-          'private'
-        elsif object_singleton_class.protected_method_defined?(@method_name)
-          'protected'
-        else
-          'public'
-        end
+        @proxy.visibility_for(@method_name)
       end
 
       # @private
       def object_singleton_class
         class << @object; self; end
@@ -96,11 +88,10 @@
       end
 
       # @private
       def restore_original_visibility
         return unless @original_visibility &&
-          (object_singleton_class.method_defined?(@method_name) ||
-           object_singleton_class.private_method_defined?(@method_name))
+          MethodReference.method_defined_at_any_visibility?(object_singleton_class, @method_name)
 
         object_singleton_class.__send__(*@original_visibility)
       end
 
       # @private