feat(message): add UnmarshalText/MarshalText methods to CompressionCodec

This allows the user to include a CompressionCodec in a structure to
be unmarshaled from a configuration file using YAML/JSON/whatever.

Author: vincentbernat

message.go

@@ -42,6 +42,28 @@ func (cc CompressionCodec) String() string {
 	}[int(cc)]
 }
 
+// UnmarshalText returns a CompressionCodec from its string representation.
+func (cc *CompressionCodec) UnmarshalText(text []byte) error {
+	codecs := map[string]CompressionCodec{
+		"none":   CompressionNone,
+		"gzip":   CompressionGZIP,
+		"snappy": CompressionSnappy,
+		"lz4":    CompressionLZ4,
+		"zstd":   CompressionZSTD,
+	}
+	codec, ok := codecs[string(text)]
+	if !ok {
+		return fmt.Errorf("cannot parse %q as a compression codec", string(text))
+	}
+	*cc = codec
+	return nil
+}
+
+// MarshalText transforms a CompressionCodec into its string representation.
+func (cc CompressionCodec) MarshalText() ([]byte, error) {
+	return []byte(cc.String()), nil
+}
+
 // Message is a kafka message type
 type Message struct {
 	Codec            CompressionCodec // codec used to compress the message contents

message_test.go

@@ -244,3 +244,32 @@ func TestMessageDecodingUnknownVersions(t *testing.T) {
 		t.Error("Decoding an unknown magic byte produced an unknown error ", err)
 	}
 }
+
+func TestCompressionCodecUnmarshal(t *testing.T) {
+	cases := []struct {
+		Input         string
+		Expected      CompressionCodec
+		ExpectedError bool
+	}{
+		{"none", CompressionNone, false},
+		{"zstd", CompressionZSTD, false},
+		{"gzip", CompressionGZIP, false},
+		{"unknown", CompressionNone, true},
+	}
+	for _, c := range cases {
+		var cc CompressionCodec
+		err := cc.UnmarshalText([]byte(c.Input))
+		if err != nil && !c.ExpectedError {
+			t.Errorf("UnmarshalText(%q) error:\n%+v", c.Input, err)
+			continue
+		}
+		if err == nil && c.ExpectedError {
+			t.Errorf("UnmarshalText(%q) got %v but expected error", c.Input, cc)
+			continue
+		}
+		if cc != c.Expected {
+			t.Errorf("UnmarshalText(%q) got %v but expected %v", c.Input, cc, c.Expected)
+			continue
+		}
+	}
+}