From 79a4d64567e6a849d88f3c819ff16a1483398f41 Mon Sep 17 00:00:00 2001
From: Vincent Bernat <vincent@bernat.ch>
Date: Mon, 7 Mar 2022 17:43:20 +0100
Subject: [PATCH] feat(message): add UnmarshalText method to CompressionCodec

This allows the user to include a CompressionCodec in a structure to
be unmarshaled from a configuration file using YAML/JSON/whatever.
---
 message.go      | 16 ++++++++++++++++
 message_test.go | 29 +++++++++++++++++++++++++++++
 2 files changed, 45 insertions(+)

diff --git a/message.go b/message.go
index fd0d1d90b7..fdc9f96beb 100644
--- a/message.go
+++ b/message.go
@@ -42,6 +42,22 @@ func (cc CompressionCodec) String() string {
 	}[int(cc)]
 }
 
+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 = CompressionCodec(codec)
+	return nil
+}
+
 // Message is a kafka message type
 type Message struct {
 	Codec            CompressionCodec // codec used to compress the message contents
diff --git a/message_test.go b/message_test.go
index a6c7cff2a5..bc4159f07a 100644
--- a/message_test.go
+++ b/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 != CompressionCodec(c.Expected) {
+			t.Errorf("UnmarshalText(%q) got %v but expected %v", c.Input, cc, c.Expected)
+			continue
+		}
+	}
+}