From a7fce75be8d6442955dc859553be403e1e3f90e6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marc-Andr=C3=A9=20Lafortune?= <github@marc-andre.ca>
Date: Wed, 18 Mar 2020 15:37:54 -0400
Subject: [PATCH] [expectations] yield_control hardening
 (rspec/rspec-expectations#1167)

Prevent invalid argument given to yield_control.exactly / at_least / at_most

Previously `exactly(:trice)` would always provide a failure with message 'expected given block to yield control', even if control was yielded.

`at_least(:trice)` yields a cryptic error 'ArgumentError: comparison of Integer with nil failedinstead of an error as it should'.

---
This commit was imported from https://github.com/rspec/rspec-expectations/commit/96bad969dc56e75b64661bc29f01207b46aaa30c.
---
 .../lib/rspec/matchers/built_in/yield.rb      | 41 +++++++++++--------
 .../rspec/matchers/built_in/yield_spec.rb     | 18 +++++++-
 2 files changed, 42 insertions(+), 17 deletions(-)

diff --git a/rspec-expectations/lib/rspec/matchers/built_in/yield.rb b/rspec-expectations/lib/rspec/matchers/built_in/yield.rb
index 929fef1ec..9879d2db1 100644
--- a/rspec-expectations/lib/rspec/matchers/built_in/yield.rb
+++ b/rspec-expectations/lib/rspec/matchers/built_in/yield.rb
@@ -98,7 +98,7 @@ def assert_valid_expect_block!
       # Not intended to be instantiated directly.
       class YieldControl < BaseMatcher
         def initialize
-          at_least(:once)
+          @expectation_type = @expected_yields_count = nil
         end
 
         # @api public
@@ -153,7 +153,7 @@ def times
         def matches?(block)
           @probe = YieldProbe.probe(block)
           return false unless @probe.has_block?
-
+          return @probe.num_yields > 0 unless @expectation_type
           @probe.num_yields.__send__(@expectation_type, @expected_yields_count)
         end
 
@@ -182,35 +182,44 @@ def supports_block_expectations?
       private
 
         def set_expected_yields_count(relativity, n)
+          raise "Multiple count constraints are not supported" if @expectation_type
+
           @expectation_type = relativity
-          @expected_yields_count = case n
-                                   when Numeric then n
-                                   when :once then 1
-                                   when :twice then 2
-                                   when :thrice then 3
-                                   end
+          @expected_yields_count = count_constraint_to_number(n)
+        end
+
+        def count_constraint_to_number(n)
+          case n
+          when Numeric then n
+          when :once then 1
+          when :twice then 2
+          when :thrice then 3
+          else
+            raise ArgumentError, "Expected a number, :once, :twice or :thrice," \
+              " but got #{n}"
+          end
         end
 
         def failure_reason
           return ' but was not a block' unless @probe.has_block?
-          return '' unless @expected_yields_count
-          " #{human_readable_expectation_type}#{human_readable_count(@expected_yields_count)}" \
-          " but yielded #{human_readable_count(@probe.num_yields)}"
+          "#{human_readable_expectation_type}#{human_readable_count(@expected_yields_count)}" \
+          " but yielded#{human_readable_count(@probe.num_yields)}"
         end
 
         def human_readable_expectation_type
           case @expectation_type
-          when :<= then 'at most '
-          when :>= then 'at least '
+          when :<= then ' at most'
+          when :>= then ' at least'
           else ''
           end
         end
 
         def human_readable_count(count)
           case count
-          when 1 then 'once'
-          when 2 then 'twice'
-          else "#{count} times"
+          when nil then ''
+          when 1 then ' once'
+          when 2 then ' twice'
+          else " #{count} times"
           end
         end
       end
diff --git a/rspec-expectations/spec/rspec/matchers/built_in/yield_spec.rb b/rspec-expectations/spec/rspec/matchers/built_in/yield_spec.rb
index 1b70c8c8f..f315ee6b5 100644
--- a/rspec-expectations/spec/rspec/matchers/built_in/yield_spec.rb
+++ b/rspec-expectations/spec/rspec/matchers/built_in/yield_spec.rb
@@ -58,7 +58,7 @@ def each_arg(*args, &block)
     it 'fails if the block does not yield' do
       expect {
         expect { |b| _dont_yield(&b) }.to yield_control
-      }.to fail_with(/expected given block to yield control/)
+      }.to fail_with(/expected given block to yield control but/)
     end
 
     it 'does not return a meaningful value from the block' do
@@ -67,6 +67,22 @@ def each_arg(*args, &block)
       expect(val).to be_nil
     end
 
+    it 'raises an error if given an invalid count argument' do
+      expect { yield_control.exactly('2') }.to raise_error(ArgumentError)
+      expect { yield_control.at_least(:trice_with_typo) }.to raise_error(ArgumentError)
+      expect { yield_control.at_most(nil) }.to raise_error(ArgumentError)
+    end
+
+    it 'is yet to support multiple calls to compatible count constraints' do
+      pending
+      expect { |b| 1.upto(4, &b) }.to yield_control.at_least(3).at_most(4).times
+      expect { |b| 1.upto(2, &b) }.not_to yield_control.at_least(3).at_most(4).times
+    end
+
+    it 'raises an error on multiple incompatible calls to count constraints' do
+      expect { yield_control.once.twice }.to raise_error(/multiple/i)
+    end
+
     context "with exact count" do
       it 'fails if the block yields wrong number of times' do
         expect {