class RSpec::Matchers::BuiltIn::YieldControl

def at_least(number)
  set_expected_yields_count(:>=, number)
  self
end

def at_most(number)
  set_expected_yields_count(:<=, number)
  self
end

def exactly(number)
  set_expected_yields_count(:==, number)
  self
end

def failure_message_for_should
  'expected given block to yield control'.tap do |failure_message|
    failure_message << relativity_failure_message
  end
end

def failure_message_for_should_not
  'expected given block not to yield control'.tap do |failure_message|
    failure_message << relativity_failure_message
  end
end

def human_readable_count
  case @expected_yields_count
  when 1 then "once"
  when 2 then "twice"
  else "#{@expected_yields_count} times"
  end
end

def human_readable_expecation_type
  case @expectation_type
  when :<= then 'at most '
  when :>= then 'at least '
  else ''
  end
end

def initialize
  @expectation_type = nil
  @expected_yields_count = nil
end

def matches?(block)
  probe = YieldProbe.probe(block)
  if @expectation_type
    probe.num_yields.send(@expectation_type, @expected_yields_count)
  else
    probe.yielded_once?(:yield_control)
  end
end

def once
  exactly(1)
  self
end

def relativity_failure_message
  return '' unless @expected_yields_count
  " #{human_readable_expecation_type}#{human_readable_count}"
end

def set_expected_yields_count(relativity, n)
  @expectation_type = relativity
  @expected_yields_count = case n
                           when Numeric then n
                           when :once then 1
                           when :twice then 2
                           end
end

def supports_block_expectations?
  true
end

def times
  self
end

def twice
  exactly(2)
  self
end