class Fluent::Test::InputTestDriver

def emit_stream(tag, es)
  @emit_streams << [tag, es.to_a]
end

def emits
  all = []
  @emit_streams.each {|tag,events|
    events.each {|time,record|
      all << [tag, time, record]
    }
  }
  all
end

def events
  all = []
  @emit_streams.each {|tag,events|
    all.concat events
  }
  all
end

def expect_emit(tag, time, record)
  (@expects ||= []) << [tag, time, record]
  self
end

def expected_emits
  @expects ||= []
end

def initialize(klass, &block)
  super(klass, &block)
  @emit_streams = []
  @expects = nil
end

def records
  all = []
  @emit_streams.each {|tag,events|
    events.each {|time,record|
      all << record
    }
  }
  all
end

def run(&block)
  m = method(:emit_stream)
  super {
    Engine.define_singleton_method(:emit_stream) {|tag,es|
      m.call(tag, es)
    }
    block.call if block
    if @expects
      i = 0
      @emit_streams.each {|tag,events|
        events.each {|time,record|
          assert_equal(@expects[i], [tag, time, record])
          i += 1
        }
      }
      assert_equal @expects.length, i
    end
  }
  self
end