Mocha #with_stubs

gistfile1.rb

=begin
block-form stubbing for Mocha; stubs get cleaned up after the block
each argument to #with_stubs should be an array of the form
  [object to stub on, method symbol, value to return]
the stub object can be an instance, a class, or Klass.any_instance
e.g. stub1 = [Foo.any_instance, :bar, bar = mock()]
     stub2 = [bar, :count, 5]
     with_stubs(stub1,stub2) { assert_equal Customer.foo.bar.count, 5  }
=end

def self.with_stubs(*params)
  stubs = params.map do |(on,method,value)|
    exp = on.stubs(method)
    exp.returns(value)
    [on, exp]
  end
  yield
ensure
  stubs.each do |(on,exp)|
    remove_expectation(on, exp)
  end
end

def self.remove_expectation(on, exp)
  mock = on.mocha
  exps = mock.__expectations__
  list = exps.to_a
  if exp && list.delete(exp)
    mm = exp.instance_variable_get(:@method_matcher)
    unless exps.matches_method?(mm.expected_method_name)
      method = on.stubba_method.new(on.stubba_object, mm.expected_method_name)
      Mocha::Mockery.instance.stubba.unstub(method)
    end
  end
end
private_class_method :remove_expectation