Reduce duplication in & between Mock & ObjectMethods

lib/mocha/mock.rb

@@ -6,7 +6,6 @@
 require 'mocha/receivers'
 require 'mocha/method_matcher'
 require 'mocha/parameters_matcher'
-require 'mocha/argument_iterator'
 require 'mocha/expectation_error_factory'
 require 'mocha/ruby_version'
 
@@ -109,14 +108,7 @@ class Mock
     #   object.expects(:expected_method_one).returns(:result_one)
     #   object.expects(:expected_method_two).returns(:result_two)
     def expects(method_name_or_hash, backtrace = nil)
-      iterator = ArgumentIterator.new(method_name_or_hash)
-      iterator.each do |*args|
-        method_name = args.shift
-        ensure_method_not_already_defined(method_name)
-        expectation = Expectation.new(self, method_name, backtrace)
-        expectation.returns(args.shift) unless args.empty?
-        @expectations.add(expectation)
-      end
+      anticipates(method_name_or_hash, backtrace)
     end
 
     # Adds an expectation that the specified method may be called any number of times with any parameters.
@@ -145,15 +137,7 @@ def expects(method_name_or_hash, backtrace = nil)
     #   object.stubs(:stubbed_method_one).returns(:result_one)
     #   object.stubs(:stubbed_method_two).returns(:result_two)
     def stubs(method_name_or_hash, backtrace = nil)
-      iterator = ArgumentIterator.new(method_name_or_hash)
-      iterator.each do |*args|
-        method_name = args.shift
-        ensure_method_not_already_defined(method_name)
-        expectation = Expectation.new(self, method_name, backtrace)
-        expectation.at_least(0)
-        expectation.returns(args.shift) unless args.empty?
-        @expectations.add(expectation)
-      end
+      anticipates(method_name_or_hash, backtrace) { |expectation| expectation.at_least(0) }
     end
 
     # Removes the specified stubbed methods (added by calls to {#expects} or {#stubs}) and all expectations associated with them.
@@ -370,5 +354,19 @@ def ensure_method_not_already_defined(method_name)
     def any_expectations?
       @expectations.any?
     end
+
+    # @private
+    def anticipates(method_name_or_hash, backtrace = nil, object = Mock.new(@mockery), &block)
+      Array(method_name_or_hash).map do |*args|
+        args = args.flatten
+        method_name = args.shift
+        Mockery.instance.stub_method(object, method_name) unless object.is_a?(Mock)
+        ensure_method_not_already_defined(method_name)
+        expectation = Expectation.new(self, method_name, backtrace)
+        expectation.returns(args.shift) unless args.empty?
+        yield expectation if block
+        @expectations.add(expectation)
+      end.last
+    end
   end
 end

lib/mocha/mockery.rb

@@ -82,6 +82,11 @@ def new_state_machine(name)
       add_state_machine(StateMachine.new(name))
     end
 
+    def stub_method(object, method_name)
+      check_stubbing(object, method_name)
+      stubba.stub(object.stubba_method_for(method_name))
+    end
+
     def verify(assertion_counter = nil)
       unless mocks.all? { |mock| mock.__verified__?(assertion_counter) }
         message = "not all expectations were satisfied\n#{mocha_inspect}"
@@ -95,7 +100,7 @@ def verify(assertion_counter = nil)
       expectations.each do |e|
         unless Mocha.configuration.stubbing_method_unnecessarily == :allow
           next if e.used?
-          on_stubbing_method_unnecessarily(e)
+          check_stubbing_method_unnecessarily(e)
         end
       end
     end
@@ -125,7 +130,15 @@ def mocha_inspect
       message
     end
 
-    def on_stubbing(object, method)
+    attr_writer :logger
+
+    def logger
+      @logger ||= Logger.new($stderr)
+    end
+
+    private
+
+    def check_stubbing(object, method)
       method = PRE_RUBY_V19 ? method.to_s : method.to_sym
       method_signature = "#{object.mocha_inspect}.#{method}"
       check(:stubbing_non_existent_method, 'non-existent method', method_signature) do
@@ -138,18 +151,10 @@ def on_stubbing(object, method)
       check(:stubbing_method_on_non_mock_object, 'method on non-mock object', method_signature)
     end
 
-    def on_stubbing_method_unnecessarily(expectation)
+    def check_stubbing_method_unnecessarily(expectation)
       check(:stubbing_method_unnecessarily, 'method unnecessarily', expectation.method_signature, expectation.backtrace)
     end
 
-    attr_writer :logger
-
-    def logger
-      @logger ||= Logger.new($stderr)
-    end
-
-    private
-
     def check(action, description, method_signature, backtrace = caller)
       return if block_given? && !yield
       message = "stubbing #{description}: #{method_signature}"

lib/mocha/object_methods.rb

@@ -1,6 +1,5 @@
 require 'mocha/mockery'
 require 'mocha/instance_method'
-require 'mocha/argument_iterator'
 require 'mocha/expectation_error_factory'
 
 module Mocha
@@ -40,6 +39,11 @@ def stubba_class
       singleton_class
     end
 
+    # @private
+    def stubba_method_for(method_name)
+      stubba_method.new(stubba_object, method_name)
+    end
+
     # Adds an expectation that the specified method must be called exactly once with any parameters.
     #
     # The original implementation of the method is replaced during the test and then restored at the end of the test. The temporary replacement method has the same visibility as the original method.
@@ -72,21 +76,7 @@ def expects(expected_methods_vs_return_values)
       if expected_methods_vs_return_values.to_s =~ /the[^a-z]*spanish[^a-z]*inquisition/i
         raise ExpectationErrorFactory.build('NOBODY EXPECTS THE SPANISH INQUISITION!')
       end
-      if frozen?
-        raise StubbingError.new("can't stub method on frozen object: #{mocha_inspect}", caller)
-      end
-      expectation = nil
-      mockery = Mocha::Mockery.instance
-      iterator = ArgumentIterator.new(expected_methods_vs_return_values)
-      iterator.each do |*args|
-        method_name = args.shift
-        mockery.on_stubbing(self, method_name)
-        method = stubba_method.new(stubba_object, method_name)
-        mockery.stubba.stub(method)
-        expectation = mocha.expects(method_name, caller)
-        expectation.returns(args.shift) unless args.empty?
-      end
-      expectation
+      anticipates(expected_methods_vs_return_values)
     end
 
     # Adds an expectation that the specified method may be called any number of times with any parameters.
@@ -118,21 +108,7 @@ def expects(expected_methods_vs_return_values)
     #
     # @see Mock#stubs
     def stubs(stubbed_methods_vs_return_values)
-      if frozen?
-        raise StubbingError.new("can't stub method on frozen object: #{mocha_inspect}", caller)
-      end
-      expectation = nil
-      mockery = Mocha::Mockery.instance
-      iterator = ArgumentIterator.new(stubbed_methods_vs_return_values)
-      iterator.each do |*args|
-        method_name = args.shift
-        mockery.on_stubbing(self, method_name)
-        method = stubba_method.new(stubba_object, method_name)
-        mockery.stubba.stub(method)
-        expectation = mocha.stubs(method_name, caller)
-        expectation.returns(args.shift) unless args.empty?
-      end
-      expectation
+      anticipates(stubbed_methods_vs_return_values) { |expectation| expectation.at_least(0) }
     end
 
     # Removes the specified stubbed methods (added by calls to {#expects} or {#stubs}) and all expectations associated with them.
@@ -161,9 +137,17 @@ def stubs(stubbed_methods_vs_return_values)
     def unstub(*method_names)
       mockery = Mocha::Mockery.instance
       method_names.each do |method_name|
-        method = stubba_method.new(stubba_object, method_name)
-        mockery.stubba.unstub(method)
+        mockery.stubba.unstub(stubba_method_for(method_name))
+      end
+    end
+
+    private
+
+    def anticipates(expected_methods_vs_return_values, &block)
+      if frozen?
+        raise StubbingError.new("can't stub method on frozen object: #{mocha_inspect}", caller)
       end
+      mocha.anticipates(expected_methods_vs_return_values, caller, self, &block)
     end
   end
 end

test/unit/mock_test.rb

@@ -49,20 +49,19 @@ def test_should_be_equal
     method_missing
     singleton_method_undefined
     initialize
+    Array
   ].freeze
 
   MACOS_EXCLUDED_METHODS =
     MACOS && MACOS_VERSION >= MACOS_MOJAVE_VERSION ? [:syscall] : []
 
-  RUBY_V19_AND_LATER_EXCLUDED_METHODS = [
-    :object_id,
-    :method_missing,
-    :singleton_method_undefined,
-    :initialize,
-    :String,
-    :singleton_method_added,
-    *MACOS_EXCLUDED_METHODS
-  ].freeze
+  RUBY_V19_AND_LATER_EXCLUDED_METHODS =
+    (PRE_RUBY_V19_EXCLUDED_METHODS.map(&:to_sym) + [
+      :object_id,
+      :String,
+      :singleton_method_added,
+      *MACOS_EXCLUDED_METHODS
+    ]).freeze
 
   OBJECT_METHODS = STANDARD_OBJECT_PUBLIC_INSTANCE_METHODS.reject do |m|
     (m =~ /^__.*__$/) ||