From b5b5c52684f7a18e329397339fcb272b5321a59a Mon Sep 17 00:00:00 2001
From: Ravi Atluri <ravi.atluri@gojek.com>
Date: Thu, 3 Aug 2023 12:51:45 +0530
Subject: [PATCH] Added FlattenMap

---
 xload/loader.go              | 14 -----------
 xload/maps.go                | 46 ++++++++++++++++++++++++++++++++++++
 xload/maps_test.go           | 32 +++++++++++++++++++++++++
 xload/providers/yaml/go.mod  |  6 ++---
 xload/providers/yaml/yaml.go | 22 ++---------------
 5 files changed, 82 insertions(+), 38 deletions(-)
 create mode 100644 xload/maps.go
 create mode 100644 xload/maps_test.go

diff --git a/xload/loader.go b/xload/loader.go
index a906bac..8e61304 100644
--- a/xload/loader.go
+++ b/xload/loader.go
@@ -25,20 +25,6 @@ func PrefixLoader(prefix string, loader Loader) Loader {
 	})
 }
 
-// MapLoader loads values from a map.
-// Mostly used for testing.
-type MapLoader map[string]string
-
-// Load fetches the value from the map.
-func (m MapLoader) Load(ctx context.Context, key string) (string, error) {
-	value, ok := m[key]
-	if !ok {
-		return "", nil
-	}
-
-	return value, nil
-}
-
 // OSLoader loads values from the OS environment.
 func OSLoader() Loader {
 	return LoaderFunc(func(ctx context.Context, key string) (string, error) {
diff --git a/xload/maps.go b/xload/maps.go
new file mode 100644
index 0000000..dfd82e3
--- /dev/null
+++ b/xload/maps.go
@@ -0,0 +1,46 @@
+package xload
+
+import (
+	"context"
+
+	"github.com/spf13/cast"
+)
+
+// MapLoader loads values from a map.
+//
+// Can be used with xload.FlattenMap as an intermediate format
+// when loading from various sources.
+type MapLoader map[string]string
+
+// Load fetches the value from the map.
+func (m MapLoader) Load(ctx context.Context, key string) (string, error) {
+	value, ok := m[key]
+	if !ok {
+		return "", nil
+	}
+
+	return value, nil
+}
+
+// FlattenMap flattens a map[string]interface{} into a map[string]string.
+// Nested maps are flattened using given separator.
+func FlattenMap(m map[string]interface{}, sep string) map[string]string {
+	return flatten(m, "", sep)
+}
+
+func flatten(m map[string]interface{}, prefix string, sep string) map[string]string {
+	flattened := make(map[string]string)
+
+	for key, value := range m {
+		switch value := value.(type) {
+		case map[string]interface{}:
+			for k, v := range flatten(value, key+sep, sep) {
+				flattened[prefix+k] = v
+			}
+		default:
+			flattened[prefix+key] = cast.ToString(value)
+		}
+	}
+
+	return flattened
+}
diff --git a/xload/maps_test.go b/xload/maps_test.go
new file mode 100644
index 0000000..90c4da9
--- /dev/null
+++ b/xload/maps_test.go
@@ -0,0 +1,32 @@
+package xload
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestFlattenMap(t *testing.T) {
+	input := map[string]interface{}{
+		"NAME":    "xload",
+		"VERSION": 1.1,
+		"AUTHOR": map[string]interface{}{
+			"NAME":  "gojek",
+			"EMAIL": "test@gojek.com",
+			"ADDRESS": map[string]interface{}{
+				"CITY": "Bombay",
+			},
+		},
+	}
+
+	want := map[string]string{
+		"NAME":                "xload",
+		"VERSION":             "1.1",
+		"AUTHOR_NAME":         "gojek",
+		"AUTHOR_EMAIL":        "test@gojek.com",
+		"AUTHOR_ADDRESS_CITY": "Bombay",
+	}
+
+	got := FlattenMap(input, "_")
+	assert.EqualValues(t, want, got)
+}
diff --git a/xload/providers/yaml/go.mod b/xload/providers/yaml/go.mod
index 2170ecf..6bd97df 100644
--- a/xload/providers/yaml/go.mod
+++ b/xload/providers/yaml/go.mod
@@ -9,12 +9,10 @@ require (
 	gopkg.in/yaml.v3 v3.0.1
 )
 
-require (
-	github.com/spf13/cast v1.5.1
-	github.com/stretchr/testify v1.8.4
-)
+require github.com/stretchr/testify v1.8.4
 
 require (
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
+	github.com/spf13/cast v1.5.1 // indirect
 )
diff --git a/xload/providers/yaml/yaml.go b/xload/providers/yaml/yaml.go
index 4b3fdfc..62f3ec3 100644
--- a/xload/providers/yaml/yaml.go
+++ b/xload/providers/yaml/yaml.go
@@ -4,14 +4,13 @@ package yaml
 import (
 	"io"
 
-	"github.com/spf13/cast"
 	"gopkg.in/yaml.v3"
 
 	"github.com/gojekfarm/xtools/xload"
 )
 
 // NewLoader reads YAML from the given io.Reader and returns a xload.Loader
-func NewLoader(r io.Reader, delim string) (xload.Loader, error) {
+func NewLoader(r io.Reader, sep string) (xload.Loader, error) {
 	b, err := io.ReadAll(r)
 	if err != nil {
 		return nil, err
@@ -24,22 +23,5 @@ func NewLoader(r io.Reader, delim string) (xload.Loader, error) {
 		return nil, err
 	}
 
-	return xload.MapLoader(flatten("", delim, out)), nil
-}
-
-func flatten(prefix, delim string, data map[string]interface{}) map[string]string {
-	flattened := make(map[string]string)
-
-	for key, value := range data {
-		switch value := value.(type) {
-		case map[string]interface{}:
-			for k, v := range flatten(key+delim, delim, value) {
-				flattened[prefix+k] = v
-			}
-		default:
-			flattened[prefix+key] = cast.ToString(value)
-		}
-	}
-
-	return flattened
+	return xload.MapLoader(xload.FlattenMap(out, sep)), nil
 }