Fix FakeKafkaConsumer to advance through messages (#221)

Author: Kelsey

kafka/kafkatest/utils.go

@@ -33,12 +33,13 @@ func KafkaPipe(log logrus.FieldLogger) (*FakeKafkaConsumer, *FakeKafkaProducer)
 }
 
 type FakeKafkaConsumer struct {
-	messages []*kafkalib.Message
-	msgMu    sync.Mutex
-	offset   int64
-	notify   chan struct{}
-	commits  chan *kafkalib.Message
-	log      logrus.FieldLogger
+	messages   []*kafkalib.Message
+	msgMu      sync.Mutex
+	offset     int64
+	readOffset int64
+	notify     chan struct{}
+	commits    chan *kafkalib.Message
+	log        logrus.FieldLogger
 }
 
 type FakeKafkaProducer struct {
@@ -63,11 +64,12 @@ func (f *FakeKafkaProducer) Close() error {
 
 func NewFakeKafkaConsumer(log logrus.FieldLogger, distri <-chan *kafkalib.Message) *FakeKafkaConsumer {
 	r := &FakeKafkaConsumer{
-		messages: make([]*kafkalib.Message, 0),
-		offset:   0,
-		notify:   make(chan struct{}),
-		log:      log,
-		commits:  make(chan *kafkalib.Message, 1000),
+		messages:   make([]*kafkalib.Message, 0),
+		offset:     0,
+		readOffset: 0,
+		notify:     make(chan struct{}),
+		log:        log,
+		commits:    make(chan *kafkalib.Message, 1000),
 	}
 
 	go func() {
@@ -86,11 +88,12 @@ func NewFakeKafkaConsumer(log logrus.FieldLogger, distri <-chan *kafkalib.Messag
 func (f *FakeKafkaConsumer) FetchMessage(ctx context.Context) (*kafkalib.Message, error) {
 	for {
 		f.msgMu.Lock()
-		if int64(len(f.messages)) > f.offset {
-			f.log.WithField("offset", f.offset).Trace("offering message")
-			msg := f.messages[f.offset]
+		if int64(len(f.messages)) > f.readOffset {
+			f.log.WithField("offset", f.readOffset).Trace("offering message")
+			msg := f.messages[f.readOffset]
 			f.msgMu.Unlock()
 
+			f.readOffset = f.readOffset + 1
 			return msg, nil
 		}
 		f.msgMu.Unlock()
@@ -121,13 +124,14 @@ func (f *FakeKafkaConsumer) CommitMessage(msg *kafkalib.Message) error {
 func (f *FakeKafkaConsumer) SetInitialOffset(offset int64) error {
 	f.msgMu.Lock()
 	f.offset = offset
+	f.readOffset = offset
 	f.msgMu.Unlock()
 	return nil
 }
 
 func (f *FakeKafkaConsumer) Seek(offset int64) error {
 	f.msgMu.Lock()
-	f.offset = offset
+	f.readOffset = offset
 	f.msgMu.Unlock()
 	return nil
 }

kafka/kafkatest/utils_test.go

@@ -15,17 +15,30 @@ func TestFakeKafkaProducer_WaitForKey(t *testing.T) {
 	defer p.Close()
 	defer c.Close()
 
-	key := []byte(`key`)
-	val := []byte(`foobar`)
-	err := p.Produce(context.Background(), &kafkalib.Message{
-		Value: val,
-		Key:   key,
+	ctx := context.Background()
+
+	err := p.Produce(ctx, &kafkalib.Message{
+		Key:   []byte(`key1`),
+		Value: []byte(`val1`),
+	})
+	require.NoError(t, err)
+
+	err = p.Produce(ctx, &kafkalib.Message{
+		Key:   []byte(`key2`),
+		Value: []byte(`val2`),
 	})
-	msg, err := c.FetchMessage(context.Background())
 	require.NoError(t, err)
-	require.Equal(t, key, msg.Key)
-	require.Equal(t, val, msg.Value)
+
+	msg, err := c.FetchMessage(ctx)
+	require.NoError(t, err)
+	require.Equal(t, "key1", string(msg.Key))
+	require.Equal(t, "val1", string(msg.Value))
+
+	msg, err = c.FetchMessage(ctx)
+	require.NoError(t, err)
+	require.Equal(t, "key2", string(msg.Key))
+	require.Equal(t, "val2", string(msg.Value))
 
 	require.NoError(t, c.CommitMessage(msg))
-	require.True(t, p.WaitForKey(key))
+	require.True(t, p.WaitForKey([]byte(`key2`)))
 }