From 94ea6cc0419e5e7c7151807619da35bbc09cb8b1 Mon Sep 17 00:00:00 2001
From: Sylvain Joyeux <sylvain.joyeux@tidewise.io>
Date: Thu, 20 Jun 2024 19:18:47 +0000
Subject: [PATCH] feat: implement explicit block support

---
 flexmock.gemspec                     |  3 +-
 lib/flexmock/argument_matchers.rb    | 15 -------
 lib/flexmock/argument_matching.rb    | 11 ++++-
 lib/flexmock/argument_types.rb       |  4 --
 lib/flexmock/call_record.rb          | 13 ++++--
 lib/flexmock/core.rb                 | 11 +++--
 lib/flexmock/expectation.rb          | 51 +++++++++++++++--------
 lib/flexmock/expectation_builder.rb  |  2 +-
 lib/flexmock/expectation_director.rb | 18 ++++----
 lib/flexmock/partial_mock.rb         |  4 +-
 lib/flexmock/validators.rb           | 35 +++-------------
 lib/flexmock/version.rb              |  2 +-
 test/assert_spy_called_test.rb       |  4 +-
 test/deprecated_methods_test.rb      |  2 +-
 test/partial_mock_test.rb            |  4 +-
 test/record_mode_test.rb             |  2 +-
 test/should_receive_test.rb          | 61 ++++++++++++++--------------
 test/spys_test.rb                    |  4 +-
 18 files changed, 115 insertions(+), 131 deletions(-)

diff --git a/flexmock.gemspec b/flexmock.gemspec
index 8c294eb..9da8e88 100644
--- a/flexmock.gemspec
+++ b/flexmock.gemspec
@@ -10,7 +10,7 @@ spec = Gem::Specification.new do |s|
     interface is simple, it is very flexible.
   }
 
-  s.required_ruby_version = ">= 2.2"
+  s.required_ruby_version = ">= 3.0"
 
   s.license = 'MIT'
 
@@ -19,7 +19,6 @@ spec = Gem::Specification.new do |s|
   s.add_development_dependency 'minitest', ">= 5.0"
   s.add_development_dependency 'rake'
   s.add_development_dependency 'simplecov', '>= 0.11.0'
-  s.add_development_dependency 'coveralls'
 
   #### Which files are to be included in this gem?  Everything!  (Except CVS directories.)
 
diff --git a/lib/flexmock/argument_matchers.rb b/lib/flexmock/argument_matchers.rb
index 0b34a67..6938372 100644
--- a/lib/flexmock/argument_matchers.rb
+++ b/lib/flexmock/argument_matchers.rb
@@ -83,19 +83,4 @@ def inspect
       "ducktype(#{@methods.map{|m| m.inspect}.join(',')})"
     end
   end
-
-  ####################################################################
-  # Match objects that implement all the methods in +methods+.
-  class OptionalProcMatcher
-    def initialize
-    end
-    def ===(target)
-      ArgumentMatching.missing?(target) || Proc === target
-    end
-    def inspect
-      "optional_proc"
-    end
-  end
-  OPTIONAL_PROC_MATCHER = OptionalProcMatcher.new
-
 end
diff --git a/lib/flexmock/argument_matching.rb b/lib/flexmock/argument_matching.rb
index 9472a0b..f32dd78 100644
--- a/lib/flexmock/argument_matching.rb
+++ b/lib/flexmock/argument_matching.rb
@@ -4,9 +4,10 @@ module ArgumentMatching
 
     MISSING_ARG = Object.new
 
-    def all_match?(expected_args, expected_kw, actual_args, actual_kw)
+    def all_match?(expected_args, expected_kw, expected_block, actual_args, actual_kw, actual_block)
       all_match_args?(expected_args, actual_args) &&
-        all_match_kw?(expected_kw, actual_kw)
+        all_match_kw?(expected_kw, actual_kw) &&
+        all_match_block?(expected_block, actual_block)
     end
 
     def all_match_args?(expected_args, actual_args)
@@ -42,6 +43,12 @@ def all_match_kw?(expected_kw, actual_kw)
       true
     end
 
+    def all_match_block?(expected_block, actual_block)
+      return true if expected_block.nil?
+
+      !(expected_block ^ actual_block)
+    end
+
     # Does the expected argument match the corresponding actual value.
     def match?(expected, actual)
       expected === actual ||
diff --git a/lib/flexmock/argument_types.rb b/lib/flexmock/argument_types.rb
index 2f82534..88223f7 100644
--- a/lib/flexmock/argument_types.rb
+++ b/lib/flexmock/argument_types.rb
@@ -47,10 +47,6 @@ def hsh(hash)
     def ducktype(*methods)
       DuckMatcher.new(methods)
     end
-
-    def optional_proc
-      OPTIONAL_PROC_MATCHER
-    end
   end
   extend ArgumentTypes
 
diff --git a/lib/flexmock/call_record.rb b/lib/flexmock/call_record.rb
index 5720d69..4422dfc 100644
--- a/lib/flexmock/call_record.rb
+++ b/lib/flexmock/call_record.rb
@@ -11,10 +11,11 @@
 
 class FlexMock
 
-  CallRecord = Struct.new(:method_name, :args, :kw, :block_given, :expectation) do
+  CallRecord = Struct.new(:method_name, :args, :kw, :block, :expectation) do
     def matches?(sym, expected_args, expected_kw, options)
       method_name == sym &&
-        ArgumentMatching.all_match?(expected_args, expected_kw, args, kw) &&
+        ArgumentMatching.all_match_args?(expected_args, args) &&
+        ArgumentMatching.all_match_kw?(expected_kw, kw) &&
         matches_block?(options[:with_block])
     end
 
@@ -22,8 +23,12 @@ def matches?(sym, expected_args, expected_kw, options)
 
     def matches_block?(block_option)
       block_option.nil? ||
-        (block_option && block_given) ||
-        (!block_option && !block_given)
+        (block_option && block) ||
+        (!block_option && !block)
+    end
+
+    def block_given
+      block
     end
   end
 
diff --git a/lib/flexmock/core.rb b/lib/flexmock/core.rb
index d58d466..93a996a 100644
--- a/lib/flexmock/core.rb
+++ b/lib/flexmock/core.rb
@@ -140,14 +140,13 @@ def by_default
   def method_missing(sym, *args, **kw, &block)
     FlexMock.verify_mocking_allowed!
 
-    enhanced_args = block_given? ? args + [block] : args
-    call_record = CallRecord.new(sym, enhanced_args, kw, block_given?)
+    call_record = CallRecord.new(sym, args, kw, block)
     @calls << call_record
     flexmock_wrap do
       if flexmock_closed?
         FlexMock.undefined
       elsif exp = flexmock_expectations_for(sym)
-        exp.call(enhanced_args, kw, call_record)
+        exp.call(args, kw, block, call_record)
       elsif @base_class && @base_class.flexmock_defined?(sym)
         FlexMock.undefined
       elsif @ignore_missing
@@ -167,9 +166,9 @@ def respond_to?(sym, *args)
   end
 
   # Find the mock expectation for method sym and arguments.
-  def flexmock_find_expectation(method_name, *args, **kw) # :nodoc:
+  def flexmock_find_expectation(method_name, *args, **kw, &block) # :nodoc:
     if exp = flexmock_expectations_for(method_name)
-      exp.find_expectation(args, kw)
+      exp.find_expectation(args, kw, block)
     end
   end
 
@@ -214,7 +213,7 @@ def flexmock_invoke_original(method_name, args, kw = {})
   # Override the built-in +method+ to include the mocked methods.
   def method(method_name)
     if (expectations = flexmock_expectations_for(method_name))
-      ->(*args, **kw) { expectations.call(args, kw) }
+      ->(*args, **kw, &block) { expectations.call(args, kw, block) }
     else
       super
     end
diff --git a/lib/flexmock/expectation.rb b/lib/flexmock/expectation.rb
index 68ce1a2..fb5ff30 100644
--- a/lib/flexmock/expectation.rb
+++ b/lib/flexmock/expectation.rb
@@ -41,6 +41,7 @@ def initialize(mock, sym, location)
       @count_validators = []
       @signature_validator = SignatureValidator.new(self)
       @count_validator_class = ExactCountValidator
+      @with_block = nil
       @actual_count = 0
       @return_value = nil
       @return_queue = []
@@ -86,48 +87,47 @@ def validate_eligible
       FlexMock.framework_adapter.check(e.message) { false }
     end
 
-    def validate_signature(args, kw)
-      @signature_validator.validate(args, kw)
+    def validate_signature(args, kw, block)
+      @signature_validator.validate(args, kw, block)
     rescue SignatureValidator::ValidationFailed => e
       FlexMock.framework_adapter.check(e.message) { false }
     end
 
     # Verify the current call with the given arguments matches the
     # expectations recorded in this object.
-    def verify_call(args, kw)
+    def verify_call(args, kw, block)
       validate_eligible
       validate_order
-      validate_signature(args, kw)
+      validate_signature(args, kw, block)
       @actual_count += 1
-      perform_yielding(args)
-      return_value(args, kw)
+      perform_yielding(block)
+      return_value(args, kw, block)
     end
 
     # Public return value (odd name to avoid accidental use as a
     # constraint).
-    def _return_value(args, kw) # :nodoc:
-      return_value(args, kw)
+    def _return_value(args, kw, block) # :nodoc:
+      return_value(args, kw, block)
     end
 
     # Find the return value for this expectation. (private version)
-    def return_value(args, kw)
+    def return_value(args, kw, block)
       case @return_queue.size
       when 0
-        block = lambda { |*a| @return_value }
+        ret_block = lambda { |*, **| @return_value }
       when 1
-        block = @return_queue.first
+        ret_block = @return_queue.first
       else
-        block = @return_queue.shift
+        ret_block = @return_queue.shift
       end
-      block.call(*args, **kw)
+      ret_block.call(*args, **kw, &block)
     end
     private :return_value
 
     # Yield stored values to any blocks given.
-    def perform_yielding(args)
+    def perform_yielding(block)
       @return_value = nil
       unless @yield_queue.empty?
-        block = args.last
         values = (@yield_queue.size == 1) ? @yield_queue.first : @yield_queue.shift
         if block && block.respond_to?(:call)
           values.each do |v|
@@ -173,8 +173,8 @@ def flexmock_verify
 
     # Does the argument list match this expectation's argument
     # specification.
-    def match_args(args, kw)
-      ArgumentMatching.all_match?(@expected_args, @expected_kw, args, kw)
+    def match_args(args, kw, block)
+      ArgumentMatching.all_match?(@expected_args, @expected_kw, @expected_block, args, kw, block)
     end
 
     # Declare that the method should expect the given argument list.
@@ -218,6 +218,23 @@ def with_kw_args(kw)
       self
     end
 
+    # Declare that the call should have a block
+    def with_block
+      @expected_block = true
+      self
+    end
+
+    # Declare that the call should have a block
+    def with_no_block
+      @expected_block = false
+      self
+    end
+
+    def with_optional_block
+      @expected_block = nil
+      self
+    end
+
     # Validate general parameters on the call signature
     def with_signature(
       required_arguments: 0, optional_arguments: 0, splat: false,
diff --git a/lib/flexmock/expectation_builder.rb b/lib/flexmock/expectation_builder.rb
index b3e17f2..c68e968 100644
--- a/lib/flexmock/expectation_builder.rb
+++ b/lib/flexmock/expectation_builder.rb
@@ -92,7 +92,7 @@ def create_demeter_chain(mock, names)
       names.each do |name|
         exp = mock.flexmock_find_expectation(name)
         if exp
-          next_mock = exp._return_value([], {})
+          next_mock = exp._return_value([], {}, nil)
           check_proper_mock(next_mock, name)
         else
           next_mock = container.flexmock("demeter_#{name}")
diff --git a/lib/flexmock/expectation_director.rb b/lib/flexmock/expectation_director.rb
index 55f0428..41873eb 100644
--- a/lib/flexmock/expectation_director.rb
+++ b/lib/flexmock/expectation_director.rb
@@ -35,8 +35,8 @@ def initialize(sym)
     # but at least we will get a good failure message).  Finally,
     # check for expectations that don't have any argument matching
     # criteria.
-    def call(args, kw, call_record=nil)
-      exp = find_expectation(args, kw)
+    def call(args, kw, block, call_record=nil)
+      exp = find_expectation(args, kw, block)
       call_record.expectation = exp if call_record
       FlexMock.check(
         proc { "no matching handler found for " +
@@ -44,7 +44,7 @@ def call(args, kw, call_record=nil)
                "\nDefined expectations:\n  " +
                @expectations.map(&:description).join("\n  ") }
       ) { !exp.nil? }
-      returned_value = exp.verify_call(args, kw)
+      returned_value = exp.verify_call(args, kw, block)
       returned_value
     end
 
@@ -54,11 +54,11 @@ def <<(expectation)
     end
 
     # Find an expectation matching the given arguments.
-    def find_expectation(args, kw) # :nodoc:
+    def find_expectation(args, kw, block) # :nodoc:
       if @expectations.empty?
-        find_expectation_in(@defaults, args, kw)
+        find_expectation_in(@defaults, args, kw, block)
       else
-        find_expectation_in(@expectations, args, kw)
+        find_expectation_in(@expectations, args, kw, block)
       end
     end
 
@@ -84,9 +84,9 @@ def defaultify_expectation(exp) # :nodoc:
 
     private
 
-    def find_expectation_in(expectations, args, kw)
-      expectations.find { |e| e.match_args(args, kw) && e.eligible? } ||
-        expectations.find { |e| e.match_args(args, kw) }
+    def find_expectation_in(expectations, args, kw, block)
+      expectations.find { |e| e.match_args(args, kw, block) && e.eligible? } ||
+        expectations.find { |e| e.match_args(args, kw, block) }
     end
   end
 
diff --git a/lib/flexmock/partial_mock.rb b/lib/flexmock/partial_mock.rb
index 06a5ad4..b8e8659 100644
--- a/lib/flexmock/partial_mock.rb
+++ b/lib/flexmock/partial_mock.rb
@@ -234,8 +234,8 @@ def flexmock_define_expectation(location, *args, **kw)
       end
     end
 
-    def flexmock_find_expectation(*args, **kw)
-      @mock.flexmock_find_expectation(*args, **kw)
+    def flexmock_find_expectation(*args, **kw, &block)
+      @mock.flexmock_find_expectation(*args, **kw, &block)
     end
 
     def add_mock_method(method_name)
diff --git a/lib/flexmock/validators.rb b/lib/flexmock/validators.rb
index a4a9bbc..a7feb76 100644
--- a/lib/flexmock/validators.rb
+++ b/lib/flexmock/validators.rb
@@ -212,44 +212,19 @@ def describe
     #
     # @param [Array] args
     # @raise ValidationFailed
-    def validate(args, kw)
-      args = args.dup
+    def validate(args, kw, block)
       kw ||= Hash.new
 
-      last_is_proc = false
-      begin
-        if args.last.kind_of?(Proc)
-          args.pop
-          last_is_proc = true
-        end
-      rescue NoMethodError
-      end
-
       if expects_keyword_arguments? && requires_keyword_arguments? && kw.empty?
         raise ValidationFailed, "#{@exp} expects keyword arguments but none were provided"
       end
 
-      # There is currently no way to disambiguate "given a block" from "given a
-      # proc as last argument" ... give some leeway in this case
-      positional_count = args.size
-
-      if required_arguments > positional_count
-        if requires_keyword_arguments?
-          raise ValidationFailed, "#{@exp} expects at least #{required_arguments} positional arguments but got only #{positional_count}"
-        end
-
-        if (required_arguments - positional_count) == 1 && last_is_proc
-          last_is_proc = false
-          positional_count += 1
-        else
-          raise ValidationFailed, "#{@exp} expects at least #{required_arguments} positional arguments but got only #{positional_count}"
-        end
+      if required_arguments > args.size
+        raise ValidationFailed, "#{@exp} expects at least #{required_arguments} positional arguments but got only #{args.size}"
       end
 
-      if !splat? && (required_arguments + optional_arguments) < positional_count
-        if !last_is_proc || (required_arguments + optional_arguments) < positional_count - 1
-          raise ValidationFailed, "#{@exp} expects at most #{required_arguments + optional_arguments} positional arguments but got #{positional_count}"
-        end
+      if !splat? && (required_arguments + optional_arguments) < args.size
+        raise ValidationFailed, "#{@exp} expects at most #{required_arguments + optional_arguments} positional arguments but got #{args.size}"
       end
 
       missing_keyword_arguments = required_keyword_arguments.
diff --git a/lib/flexmock/version.rb b/lib/flexmock/version.rb
index f78f34c..5f3e4e7 100644
--- a/lib/flexmock/version.rb
+++ b/lib/flexmock/version.rb
@@ -1,3 +1,3 @@
 class FlexMock
-    VERSION = "2.3.8"
+    VERSION = "3.0.0"
 end
diff --git a/test/assert_spy_called_test.rb b/test/assert_spy_called_test.rb
index 6b07d43..f9e2004 100644
--- a/test/assert_spy_called_test.rb
+++ b/test/assert_spy_called_test.rb
@@ -57,8 +57,8 @@ def test_assert_rejects_incorrect_type
   def test_assert_detects_blocks
     spy.foo { }
     spy.bar
-    assert_spy_called spy, :foo, Proc
-    assert_spy_called spy, :bar
+    assert_spy_called spy, {with_block: true}, :foo
+    assert_spy_called spy, {with_block: false}, :bar
   end
 
   def test_assert_detects_any_args
diff --git a/test/deprecated_methods_test.rb b/test/deprecated_methods_test.rb
index 59da1e8..796e47e 100644
--- a/test/deprecated_methods_test.rb
+++ b/test/deprecated_methods_test.rb
@@ -39,7 +39,7 @@ def test_handle_no_block
 
   def test_called_with_block
     called = false
-    s { @mock.mock_handle(:blip) { |block| block.call } }
+    s { @mock.mock_handle(:blip) { |&block| block.call } }
     @mock.blip { called = true }
     assert called, "Block to blip should be called"
   end
diff --git a/test/partial_mock_test.rb b/test/partial_mock_test.rb
index 4b49654..1f86752 100644
--- a/test/partial_mock_test.rb
+++ b/test/partial_mock_test.rb
@@ -74,8 +74,8 @@ def test_stub_command_can_configure_via_block
 
   def test_stubbed_methods_can_take_blocks
     obj = Object.new
-    flexmock(obj).should_receive(:with_block).once.with(Proc).
-      and_return { |block| block.call }
+    flexmock(obj).should_receive(:with_block).once.with_block.
+      and_return { |&block| block.call }
     assert_equal :block, obj.with_block { :block }
   end
 
diff --git a/test/record_mode_test.rb b/test/record_mode_test.rb
index 0622a26..23d5975 100644
--- a/test/record_mode_test.rb
+++ b/test/record_mode_test.rb
@@ -61,7 +61,7 @@ def test_recording_mode_does_not_specify_order
   def test_recording_mode_gets_block_args_too
     mock = flexmock("mock")
     mock.should_expect do |r|
-      r.f(1, Proc) { |arg, block|
+      r.f(1) { |arg, &block|
         refute_nil block
         block.call
       }
diff --git a/test/should_receive_test.rb b/test/should_receive_test.rb
index 1c19bbf..31dd0b5 100644
--- a/test/should_receive_test.rb
+++ b/test/should_receive_test.rb
@@ -88,9 +88,9 @@ def test_returns_with_block
     end
   end
 
-  def test_block_example_from_readme
+  def test_with_block_example_from_readme
     FlexMock.use do |m|
-      m.should_receive(:foo).with(Integer,Proc).and_return(:got_block)
+      m.should_receive(:foo).with(Integer).with_block.and_return(:got_block)
       m.should_receive(:foo).with(Integer).and_return(:no_block)
 
       assert_equal :no_block, m.foo(1)
@@ -98,6 +98,25 @@ def test_block_example_from_readme
     end
   end
 
+  def test_with_no_block_example_from_readme
+    FlexMock.use do |m|
+      m.should_receive(:foo).with(Integer).with_no_block.and_return(:no_block)
+      m.should_receive(:foo).with(Integer).and_return(:got_block)
+
+      assert_equal :no_block, m.foo(1)
+      assert_equal :got_block, m.foo(1) { }
+    end
+  end
+
+  def test_with_optional_block
+    FlexMock.use do |m|
+      m.should_receive(:foo).with(Integer).with_optional_block.twice
+
+      m.foo(1)
+      m.foo(1) {}
+    end
+  end
+
   def test_return_with_and_without_block_interleaved
     FlexMock.use do |m|
       m.should_receive(:hi).and_return(:a).and_return { :b }.and_return(:c)
@@ -467,28 +486,6 @@ def test_with_equal_arg_nonmatching
     end
   end
 
-  def test_with_optional_proc
-    FlexMock.use('greeter') do |m|
-      m.should_receive(:hi).with(optional_proc).once
-      m.hi { }
-    end
-  end
-
-  def test_with_optional_proc_and_missing_proc
-    FlexMock.use('greeter') do |m|
-      m.should_receive(:hi).with(optional_proc).once
-      m.hi
-    end
-  end
-
-  def test_with_optional_proc_distinquishes_between_nil_and_missing
-    FlexMock.use('greeter') do |m|
-      m.should_receive(:hi).with(optional_proc).never
-      m.should_receive(:hi).with(nil).once
-      m.hi(nil)
-    end
-  end
-
   def test_with_arbitrary_arg_matching
     FlexMock.use('greeter') do |m|
       m.should_receive(:hi).with(FlexMock.on { |arg| arg % 2 == 0 rescue nil }).twice
@@ -552,9 +549,9 @@ def test_arg_matching_with_string_doesnt_over_match
     end
   end
 
-  def test_block_arg_given_to_no_args
+  def test_block_arg_given_to_no_block
     FlexMock.use do |m|
-      m.should_receive(:hi).with_no_args.returns(20)
+      m.should_receive(:hi).with_no_block.returns(20)
       assert_mock_failure(check_failed_error, :message =>NO_MATCH_ERROR_MESSAGE, :deep => true, :line => __LINE__+1) {
         m.hi { 1 }
       }
@@ -564,8 +561,9 @@ def test_block_arg_given_to_no_args
   def test_block_arg_given_to_matching_proc
     FlexMock.use do |m|
       arg = nil
-      m.should_receive(:hi).with(Proc).once.
-        and_return { |block| arg = block; block.call }
+      m.should_receive(:hi)
+       .with_block.once
+       .and_return { |&block| arg = block; block.call }
       result = m.hi { 1 }
       assert_equal 1, arg.call
       assert_equal 1, result
@@ -1229,11 +1227,14 @@ def test_with_signature_required_arguments
     end
   end
 
-  def test_a_proc_argument_last_can_be_interpreted_as_positional_argument
+  def test_a_proc_argument_last_is_not_interpreted_as_positional_argument
     FlexMock.use do |mock|
       mock.should_receive(:m).with_signature(required_arguments: 2)
-      mock.m(1) { }
       mock.m(1, 2) { }
+
+      assert_raises(FlexMock::CheckFailedError) do
+        mock.m(1) { }
+      end
     end
   end
 
diff --git a/test/spys_test.rb b/test/spys_test.rb
index 8dbd883..fe6c54a 100644
--- a/test/spys_test.rb
+++ b/test/spys_test.rb
@@ -67,7 +67,7 @@ def test_spy_rejects_if_times_options_not_matching
 
   def test_spy_detects_a_block
     @spy.foo { }
-    assert_spy_called @spy, :foo, Proc
+    assert_spy_called @spy, {:with_block => true}, :foo
   end
 
   def test_spy_rejects_a_block
@@ -87,7 +87,7 @@ def test_spy_rejects_a_missing_block
 
   def test_spy_ignores_block
     @spy.foo { }
-    assert_spy_called @spy, :foo, Proc
+    assert_spy_called @spy, :foo
   end
 
   def test_spy_accepts_correct_additional_validations