Making enumerator from method that yields values

Original challenge related to AWS SQS QueuePoller

The challenge was to test a static method that yields multiple values but should stop, when some condition is met.

Let’s imagine such method as follows

class Poller
  def self.poll(condition = -> {})
    counter = 0

    while true do
      yield counter += 1
      break if condition.call
    end
  end
end

The problem is with testing such method. We do not only need to test what it yields, but we also need to test and control when it stops. To control when it stops, we need to access the actual block, but to test what it yields, we either need yield_successive_args matcher or we need to fetch consecutive results.

It is possible by aggregating each yielded value and then asserting them altogether, but the resultant code is not nice. The solution would be to make an Enumerator from the poll method and use next to get consecutive results. It is also easy as described in this blog post. The problem is, that we do not want to write code that is only required by our tests.

So the idea is to add creating enumerators when the block is not provided dynamically to the class when testing.

Poller.define_singleton_method(:poll_with_enum) do |*args, &block| 
    return enum_for(:poll) unless block.present?
    poll_without_enum(*args, &block)
end
# alias_method_chain is deprecated
# Poller.singleton_class.alias_method_chain(:poll, :enum)
Poller.singleton_class.alias_method :poll_without_enum, :poll
Poller.singleton_class.alias_method :poll, :poll_with_enum

if we turn this into a helper…

def with_enumerated(subject, method_name)
  begin
    subject.define_singleton_method("#{method_name}_with_enum") do |*args, &block|
      return enum_for(method_name, *args) unless block.present?
      public_send("#{method_name}_without_enum",*args, &block)
    end

    subject.singleton_class.alias_method "#{method_name}_without_enum", method_name
    subject.singleton_class.alias_method method_name, "#{method_name}_with_enum"

    yield
  ensure
    subject.singleton_class.alias_method method_name, "#{method_name}_without_enum"
    subject.singleton_class.remove_method "#{method_name}_with_enum"
  end
end

…then we could leverage it in our tests!

with_enumerated(Poller, :poll) do
  $stop = false
  poller = Poller.poll(condition = -> { $stop == true })

  first_value = poller.next
  expect(first_value).to eq 1

  $stop = true
  second_value = poller.next
  expect(second_value).to eq 2

  expect { poller.next }.to raise_exception(StopIteration)
end

Tweet