config: Support converting string configs to []byte (flyteorg#85)

* config: Support converting string configs to []byte

Signed-off-by: Haytham Abuelfutuh <[email protected]>

* Use base64 encoding to match yaml.unmarshal behavior

Signed-off-by: Haytham Abuelfutuh <[email protected]>

* support string array since it seems mapstructure does that

Signed-off-by: Haytham Abuelfutuh <[email protected]>

* Fix unit test value

Signed-off-by: Haytham Abuelfutuh <[email protected]>

* Increase patch coverage

Signed-off-by: Haytham Abuelfutuh <[email protected]>

flytestdlib/config/tests/testdata/config.yaml

@@ -9,3 +9,4 @@ other-component:
     - world
     - '!'
   url-value: http://something.com
+  myByteArray: JDJhJDA2JHB4czFBa0c4MUt2cmhwbWwxUWlMU09RYVRrOWVlUHJVLzdZYWI5eTA3aDN4MFRnbGJhb1Q2

flytestdlib/config/tests/types_test.go

@@ -25,6 +25,7 @@ type OtherComponentConfig struct {
 	IntValue                int             `json:"int-val"`
 	StringArray             []string        `json:"strings"`
 	StringArrayWithDefaults []string        `json:"strings-def"`
+	MyByteArray             []byte          `json:"myByteArray"`
 }
 
 type Item struct {

flytestdlib/config/viper/viper.go

@@ -2,6 +2,7 @@ package viper
 
 import (
 	"context"
+	"encoding/base64"
 	"encoding/json"
 	"flag"
 	"fmt"
@@ -193,6 +194,35 @@ func sliceToMapHook(f reflect.Kind, t reflect.Kind, data interface{}) (interface
 	return data, nil
 }
 
+// stringToByteArray allows the conversion from strings to []byte. mapstructure's default behavior involve converting
+// each element as a uint8 before assembling the final []byte.
+func stringToByteArray(f, t reflect.Type, data interface{}) (interface{}, error) {
+	// Only handle string -> []byte conversion
+	if t.Kind() != reflect.Slice || t.Elem().Kind() != reflect.Uint8 {
+		return data, nil
+	}
+
+	asStr := ""
+	if f.Kind() == reflect.String {
+		asStr = data.(string)
+	} else if f.Kind() == reflect.Slice && f.Elem().Kind() == reflect.String {
+		asSlice := data.([]string)
+		if len(asSlice) == 0 {
+			return data, nil
+		}
+
+		asStr = asSlice[0]
+	}
+
+	b := make([]byte, base64.StdEncoding.DecodedLen(len(asStr)))
+	n, err := base64.StdEncoding.Decode(b, []byte(asStr))
+	if err != nil {
+		return nil, err
+	}
+
+	return b[:n], nil
+}
+
 // This decoder hook tests types for json unmarshaling capability. If implemented, it uses json unmarshal to build the
 // object. Otherwise, it'll just pass on the original data.
 func jsonUnmarshallerHook(_, to reflect.Type, data interface{}) (interface{}, error) {
@@ -298,6 +328,7 @@ func defaultDecoderConfig(output interface{}, opts ...viperLib.DecoderConfigOpti
 			mapstructure.StringToTimeDurationHookFunc(),
 			mapstructure.StringToSliceHookFunc(","),
 			sliceToMapHook,
+			stringToByteArray,
 		),
 		// Empty/zero fields before applying provided values. This avoids potentially undesired/unexpected merging logic.
 		ZeroFields: true,

flytestdlib/config/viper/viper_test.go

@@ -0,0 +1,54 @@
+package viper
+
+import (
+	"encoding/base64"
+	"reflect"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_stringToByteArray(t *testing.T) {
+	t.Run("Expected types", func(t *testing.T) {
+		input := "hello world"
+		base64Encoded := base64.StdEncoding.EncodeToString([]byte(input))
+		res, err := stringToByteArray(reflect.TypeOf(base64Encoded), reflect.TypeOf([]byte{}), base64Encoded)
+		assert.NoError(t, err)
+		assert.Equal(t, []byte(input), res)
+	})
+
+	t.Run("Expected types - array string", func(t *testing.T) {
+		input := []string{"hello world"}
+		base64Encoded := base64.StdEncoding.EncodeToString([]byte(input[0]))
+		res, err := stringToByteArray(reflect.TypeOf(input), reflect.TypeOf([]byte{}), []string{base64Encoded})
+		assert.NoError(t, err)
+		assert.Equal(t, []byte(input[0]), res)
+	})
+
+	t.Run("Expected types - invalid encoding", func(t *testing.T) {
+		input := []string{"hello world"}
+		_, err := stringToByteArray(reflect.TypeOf(input), reflect.TypeOf([]byte{}), []string{"invalid base64"})
+		assert.Error(t, err)
+	})
+
+	t.Run("Expected types - empty array string", func(t *testing.T) {
+		input := []string{"hello world"}
+		res, err := stringToByteArray(reflect.TypeOf(input), reflect.TypeOf([]byte{}), []string{})
+		assert.NoError(t, err)
+		assert.Equal(t, []string{}, res)
+	})
+
+	t.Run("Unexpected types", func(t *testing.T) {
+		input := 5
+		res, err := stringToByteArray(reflect.TypeOf(input), reflect.TypeOf([]byte{}), input)
+		assert.NoError(t, err)
+		assert.NotEqual(t, []byte("hello"), res)
+	})
+
+	t.Run("Unexpected types", func(t *testing.T) {
+		input := 5
+		res, err := stringToByteArray(reflect.TypeOf(input), reflect.TypeOf(""), input)
+		assert.NoError(t, err)
+		assert.NotEqual(t, []byte("hello"), res)
+	})
+}