add test for batching messages with headers

Author: dominicbarnes

writer_test.go

@@ -7,6 +7,7 @@ import (
 	"io"
 	"math"
 	"strconv"
+	"strings"
 	"sync"
 	"testing"
 	"time"
@@ -134,6 +135,10 @@ func TestWriter(t *testing.T) {
 			scenario: "writing messages with a small batch byte size",
 			function: testWriterSmallBatchBytes,
 		},
+		{
+			scenario: "writing messages with headers",
+			function: testWriterBatchBytesHeaders,
+		},
 		{
 			scenario: "setting a non default balancer on the writer",
 			function: testWriterSetsRightBalancer,
@@ -592,6 +597,67 @@ func testWriterSmallBatchBytes(t *testing.T) {
 	}
 }
 
+func testWriterBatchBytesHeaders(t *testing.T) {
+	topic := makeTopic()
+	createTopic(t, topic, 1)
+	defer deleteTopic(t, topic)
+
+	offset, err := readOffset(topic, 0)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	w := newTestWriter(WriterConfig{
+		Topic:        topic,
+		BatchBytes:   100,
+		BatchTimeout: 50 * time.Millisecond,
+		Balancer:     &RoundRobin{},
+	})
+	defer w.Close()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	if err := w.WriteMessages(ctx, []Message{
+		{
+			Value: []byte("Hello World 1"),
+			Headers: []Header{
+				{Key: "User-Agent", Value: []byte("abc/xyz")},
+			},
+		},
+		{
+			Value: []byte("Hello World 2"),
+			Headers: []Header{
+				{Key: "User-Agent", Value: []byte("abc/xyz")},
+			},
+		},
+	}...); err != nil {
+		t.Error(err)
+		return
+	}
+	ws := w.Stats()
+	if ws.Writes != 2 {
+		t.Error("didn't batch messages; Writes: ", ws.Writes)
+		return
+	}
+	msgs, err := readPartition(topic, 0, offset)
+	if err != nil {
+		t.Error("error reading partition", err)
+		return
+	}
+
+	if len(msgs) != 2 {
+		t.Error("bad messages in partition", msgs)
+		return
+	}
+
+	for _, m := range msgs {
+		if strings.HasPrefix(string(m.Value), "Hello World") {
+			continue
+		}
+		t.Error("bad messages in partition", msgs)
+	}
+}
+
 func testWriterMultipleTopics(t *testing.T) {
 	topic1 := makeTopic()
 	createTopic(t, topic1, 1)