Function TestConsumerRebalancingMultiplePartitions

consumer_test.go:1224–1422

If leadership for a partition is changing then consumer resolves the new leader and switches to it.

(t *testing.T)

// If leadership for a partition is changing then consumer resolves the new
// leader and switches to it.
func TestConsumerRebalancingMultiplePartitions(t *testing.T) {
	// initial setup
	seedBroker := NewMockBroker(t, 10)
	leader0 := NewMockBroker(t, 0)
	leader1 := NewMockBroker(t, 1)

	seedBroker.SetHandlerByMap(map[string]MockResponse{
		"MetadataRequest": NewMockMetadataResponse(t).
			SetBroker(leader0.Addr(), leader0.BrokerID()).
			SetBroker(leader1.Addr(), leader1.BrokerID()).
			SetBroker(seedBroker.Addr(), seedBroker.BrokerID()).
			SetLeader("my_topic", 0, leader0.BrokerID()).
			SetLeader("my_topic", 1, leader1.BrokerID()),
	})

	mockOffsetResponse1 := NewMockOffsetResponse(t).
		SetOffset("my_topic", 0, OffsetOldest, 0).
		SetOffset("my_topic", 0, OffsetNewest, 1000).
		SetOffset("my_topic", 1, OffsetOldest, 0).
		SetOffset("my_topic", 1, OffsetNewest, 1000)
	leader0.SetHandlerByMap(map[string]MockResponse{
		"OffsetRequest": mockOffsetResponse1,
		"FetchRequest":  NewMockFetchResponse(t, 1),
	})
	leader1.SetHandlerByMap(map[string]MockResponse{
		"OffsetRequest": mockOffsetResponse1,
		"FetchRequest":  NewMockFetchResponse(t, 1),
	})

	// launch test goroutines
	config := NewTestConfig()
	config.ClientID = t.Name()
	config.Consumer.Retry.Backoff = 50
	master, err := NewConsumer([]string{seedBroker.Addr()}, config)
	if err != nil {
		t.Fatal(err)
	}

	consumers := map[int32]PartitionConsumer{}
	checkMessage := func(partition int32, offset int) {
		c := consumers[partition]
		message := <-c.Messages()
		t.Logf("Received message my_topic-%d offset=%d", partition, message.Offset)
		if message.Offset != int64(offset) {
			t.Error("Incorrect message offset!", offset, partition, message.Offset)
		}
		if message.Partition != partition {
			t.Error("Incorrect message partition!")
		}
	}

	for i := range int32(2) {
		consumer, err := master.ConsumePartition("my_topic", i, 0)
		if err != nil {
			t.Fatal(err)
		}

		go func(c PartitionConsumer) {