func mockNewBroker(t *testing.T, cp ChainPartition) (Broker, error) {
mockBroker := sarama.NewMockBroker(t, testBrokerID)
handlerMap := make(map[string]sarama.MockResponse)
// The sarama mock package doesn't allow us to return an error
// for invalid offset requests, so we return an offset of -1.
// Note that the mock offset responses below imply a broker with
// newestOffset-1 blocks available. Therefore, if you are using this
// broker as part of a bigger test where you intend to consume blocks,
// make sure that the mockConsumer has been initialized accordingly
// (Set the 'offset' parameter to newestOffset-1.)
handlerMap["OffsetRequest"] = sarama.NewMockOffsetResponse(t).
SetOffset(cp.Topic(), cp.Partition(), sarama.OffsetOldest, testOldestOffset).
SetOffset(cp.Topic(), cp.Partition(), sarama.OffsetNewest, testNewestOffset)
mockBroker.SetHandlerByMap(handlerMap)
broker := sarama.NewBroker(mockBroker.Addr())
if err := broker.Open(nil); err != nil {
return nil, fmt.Errorf("Cannot connect to mock broker: %s", err)
}
return &mockBrockerImpl{
brokerImpl: brokerImpl{
broker: broker,
},
mockBroker: mockBroker,
handlerMap: handlerMap,
}, nil
}
func TestNewBrokerReturnsPartitionLeader(t *testing.T) {
cp := newChainPartition(provisional.TestChainID, rawPartition)
broker1 := sarama.NewMockBroker(t, 1)
broker2 := sarama.NewMockBroker(t, 2)
broker3 := sarama.NewMockBroker(t, 3)
defer func() {
broker2.Close()
broker3.Close()
}()
// Use broker1 and broker2 as bootstrap brokers, but shutdown broker1 right away
broker1.Close()
// Add expectation that broker2 will return a metadata response
// that identifies broker3 as the topic partition leader
broker2.SetHandlerByMap(map[string]sarama.MockResponse{
"MetadataRequest": sarama.NewMockMetadataResponse(t).
SetBroker(broker1.Addr(), broker1.BrokerID()).
SetBroker(broker2.Addr(), broker2.BrokerID()).
SetBroker(broker3.Addr(), broker3.BrokerID()).
SetLeader(cp.Topic(), cp.Partition(), broker3.BrokerID()),
})
// Add expectation that broker3 responds to an offset request
broker3.SetHandlerByMap(map[string]sarama.MockResponse{
"OffsetRequest": sarama.NewMockOffsetResponse(t).
SetOffset(cp.Topic(), cp.Partition(), sarama.OffsetOldest, testOldestOffset).
SetOffset(cp.Topic(), cp.Partition(), sarama.OffsetNewest, testNewestOffset),
})
// Get leader for the test chain partition
leaderBroker, _ := newBroker([]string{broker1.Addr(), broker2.Addr()}, cp)
// Only broker3 will respond successfully to an offset request
offsetRequest := new(sarama.OffsetRequest)
offsetRequest.AddBlock(cp.Topic(), cp.Partition(), -1, 1)
if _, err := leaderBroker.GetOffset(cp, offsetRequest); err != nil {
t.Fatal("Expected leader broker to respond to request:", err)
}
}
func TestReceivePayloadMessage(t *testing.T) {
broker := sarama.NewMockBroker(t, 1)
ctrl := gomock.NewController(t)
tmpDir, tmpErr := ioutil.TempDir("", "kafkainput-tests")
if tmpErr != nil {
t.Errorf("Unable to create a temporary directory: %s", tmpErr)
}
defer func() {
if err := os.RemoveAll(tmpDir); err != nil {
t.Errorf("Cleanup failed: %s", err)
}
ctrl.Finish()
}()
topic := "test"
mockFetchResponse := sarama.NewMockFetchResponse(t, 1)
mockFetchResponse.SetMessage(topic, 0, 0, sarama.ByteEncoder([]byte{0x41, 0x42}))
broker.SetHandlerByMap(map[string]sarama.MockResponse{
"MetadataRequest": sarama.NewMockMetadataResponse(t).
SetBroker(broker.Addr(), broker.BrokerID()).
SetLeader(topic, 0, broker.BrokerID()),
"OffsetRequest": sarama.NewMockOffsetResponse(t).
SetOffset(topic, 0, sarama.OffsetOldest, 0).
SetOffset(topic, 0, sarama.OffsetNewest, 2),
"FetchRequest": mockFetchResponse,
})
pConfig := NewPipelineConfig(nil)
pConfig.Globals.BaseDir = tmpDir
ki := new(KafkaInput)
ki.SetName(topic)
ki.SetPipelineConfig(pConfig)
config := ki.ConfigStruct().(*KafkaInputConfig)
config.Addrs = append(config.Addrs, broker.Addr())
config.Topic = topic
ith := new(plugins_ts.InputTestHelper)
ith.Pack = NewPipelinePack(pConfig.InputRecycleChan())
ith.MockHelper = pipelinemock.NewMockPluginHelper(ctrl)
ith.MockInputRunner = pipelinemock.NewMockInputRunner(ctrl)
ith.MockSplitterRunner = pipelinemock.NewMockSplitterRunner(ctrl)
err := ki.Init(config)
if err != nil {
t.Fatalf("%s", err)
}
ith.MockInputRunner.EXPECT().NewSplitterRunner("").Return(ith.MockSplitterRunner)
ith.MockSplitterRunner.EXPECT().UseMsgBytes().Return(false)
ith.MockSplitterRunner.EXPECT().Done()
decChan := make(chan func(*PipelinePack), 1)
decCall := ith.MockSplitterRunner.EXPECT().SetPackDecorator(gomock.Any())
decCall.Do(func(dec func(pack *PipelinePack)) {
decChan <- dec
})
bytesChan := make(chan []byte, 1)
splitCall := ith.MockSplitterRunner.EXPECT().SplitBytes(gomock.Any(), nil)
splitCall.Do(func(recd []byte, del Deliverer) {
bytesChan <- recd
})
errChan := make(chan error)
go func() {
errChan <- ki.Run(ith.MockInputRunner, ith.MockHelper)
}()
recd := <-bytesChan
if string(recd) != "AB" {
t.Errorf("Invalid Payload Expected: AB received: %s", string(recd))
}
packDec := <-decChan
packDec(ith.Pack)
if ith.Pack.Message.GetType() != "heka.kafka" {
t.Errorf("Invalid Type %s", ith.Pack.Message.GetType())
}
// There is a hang on the consumer close with the mock broker
// closing the brokers before the consumer works around the issue
// and is good enough for this test.
broker.Close()
ki.Stop()
err = <-errChan
if err != nil {
t.Fatal(err)
}
filename := filepath.Join(tmpDir, "kafka", "test.test.0.offset.bin")
if o, err := readCheckpoint(filename); err != nil {
t.Errorf("Could not read the checkpoint file: %s", filename)
} else {
if o != 1 {
t.Errorf("Incorrect offset Expected: 1 Received: %d", o)
}
}
}