From 68b69565b9d8cad93227870d2a42e598ce242b85 Mon Sep 17 00:00:00 2001
From: Collin Van Dyck <collin@segment.com>
Date: Tue, 5 Oct 2021 11:28:50 -0400
Subject: [PATCH] Add support for testing message headers.

For v2 message set builders, we serialize message headers, and
then test that the headers are set on the messages read from
the message set reader.
---
 builder_test.go |  8 +++++++-
 message_test.go | 28 +++++++++++++++++++++++-----
 2 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/builder_test.go b/builder_test.go
index e1c4ba9d4..6f9b16b6e 100644
--- a/builder_test.go
+++ b/builder_test.go
@@ -187,7 +187,13 @@ func (f v2MessageSetBuilder) bytes() []byte {
 							wb.Write(msg.Key)                                             // key bytes
 							wb.writeVarInt(int64(len(msg.Value)))                         // value len
 							wb.Write(msg.Value)                                           // value bytes
-							wb.writeVarInt(0)                                             // number of headers
+							wb.writeVarInt(int64(len(msg.Headers)))                       // number of headers
+							for _, header := range msg.Headers {
+								wb.writeVarInt(int64(len(header.Key)))
+								wb.Write([]byte(header.Key))
+								wb.writeVarInt(int64(len(header.Value)))
+								wb.Write(header.Value)
+							}
 						})
 						wb.writeVarInt(int64(len(bs)))
 						wb.Write(bs)
diff --git a/message_test.go b/message_test.go
index 08477f89f..9d4a18873 100644
--- a/message_test.go
+++ b/message_test.go
@@ -29,6 +29,12 @@ func TestMessageSetReader(t *testing.T) {
 			Offset: int64(i + startOffset),
 			Key:    []byte(fmt.Sprintf("key-%d", i)),
 			Value:  []byte(fmt.Sprintf("val-%d", i)),
+			Headers: []Header{
+				{
+					Key:   fmt.Sprintf("header-key-%d", i),
+					Value: []byte(fmt.Sprintf("header-value-%d", i)),
+				},
+			},
 		}
 	}
 	defaultHeader := fetchResponseHeader{
@@ -280,11 +286,23 @@ func TestMessageSetReader(t *testing.T) {
 			}
 			rh.offset = tc.builder.messages()[0].Offset
 			rh.debug = tc.debug
-			for _, expected := range tc.builder.messages() {
-				msg := rh.readMessage()
-				require.Equal(t, string(expected.Key), string(msg.Key))
-				require.Equal(t, string(expected.Value), string(msg.Value))
-				require.Equal(t, expected.Offset, msg.Offset)
+			for _, messageSet := range tc.builder.msgSets {
+				for _, expected := range messageSet.messages() {
+					msg := rh.readMessage()
+					require.Equal(t, string(expected.Key), string(msg.Key))
+					require.Equal(t, string(expected.Value), string(msg.Value))
+					switch messageSet.(type) {
+					case v0MessageSetBuilder, v1MessageSetBuilder:
+						// v0 and v1 message sets do not have headers
+						require.Len(t, msg.Headers, 0)
+					case v2MessageSetBuilder:
+						// v2 message sets can have headers
+						require.EqualValues(t, expected.Headers, msg.Headers)
+					default:
+						t.Fatalf("unknown builder: %T", messageSet)
+					}
+					require.Equal(t, expected.Offset, msg.Offset)
+				}
 			}
 			// verify the reader stack is empty
 			require.EqualValues(t, 0, rh.remain)