From 285f1510193473600929e1183915f6220fabab94 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Masson?=
 <benoitmasson@users.noreply.github.com>
Date: Sun, 23 Oct 2016 23:04:21 +0200
Subject: [PATCH] Fixed AllKeys() to include env values added with BindEnv()

* Fixed: values bound with BindEnv added to AllKeys()

Cast was not working, and v.env wasn't used when merging keys.

Rewrote explicit and specific casts for maps storing strings or FlagValues.

* Added: test for BindEnv() and AllKeys()

To make sure AllSettings() and Unmarshal() will consider environment
variables added with BindEnv().
---
 viper.go      | 23 +++++++++++------------
 viper_test.go | 17 +++++++++++++++++
 2 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/viper.go b/viper.go
index 8d35fda..dba229b 100644
--- a/viper.go
+++ b/viper.go
@@ -1162,6 +1162,14 @@ func castMapStringToMapInterface(src map[string]string) map[string]interface{} {
 	return tgt
 }
 
+func castMapFlagToMapInterface(src map[string]FlagValue) map[string]interface{} {
+	tgt := map[string]interface{}{}
+	for k, v := range src {
+		tgt[k] = v
+	}
+	return tgt
+}
+
 // mergeMaps merges two maps. The `itgt` parameter is for handling go-yaml's
 // insistence on parsing nested structures as `map[interface{}]interface{}`
 // instead of using a `string` as the key for nest structures beyond one level
@@ -1307,8 +1315,8 @@ func (v *Viper) AllKeys() []string {
 	// add all paths, by order of descending priority to ensure correct shadowing
 	m = v.flattenAndMergeMap(m, castMapStringToMapInterface(v.aliases), "")
 	m = v.flattenAndMergeMap(m, v.override, "")
-	m = v.mergeFlatMap(m, v.pflags)
-	m = v.mergeFlatMap(m, v.env)
+	m = v.mergeFlatMap(m, castMapFlagToMapInterface(v.pflags))
+	m = v.mergeFlatMap(m, castMapStringToMapInterface(v.env))
 	m = v.flattenAndMergeMap(m, v.config, "")
 	m = v.flattenAndMergeMap(m, v.kvstore, "")
 	m = v.flattenAndMergeMap(m, v.defaults, "")
@@ -1360,16 +1368,7 @@ func (v *Viper) flattenAndMergeMap(shadow map[string]bool, m map[string]interfac
 
 // mergeFlatMap merges the given maps, excluding values of the second map
 // shadowed by values from the first map.
-func (v *Viper) mergeFlatMap(shadow map[string]bool, mi interface{}) map[string]bool {
-	// unify input map
-	var m map[string]interface{}
-	switch mi.(type) {
-	case map[string]string, map[string]FlagValue:
-		m = cast.ToStringMap(mi)
-	default:
-		return shadow
-	}
-
+func (v *Viper) mergeFlatMap(shadow map[string]bool, m map[string]interface{}) map[string]bool {
 	// scan keys
 outer:
 	for k, _ := range m {
diff --git a/viper_test.go b/viper_test.go
index ab1f834..83dbac4 100644
--- a/viper_test.go
+++ b/viper_test.go
@@ -445,6 +445,23 @@ func TestAllKeys(t *testing.T) {
 	assert.Equal(t, all, AllSettings())
 }
 
+func TestAllKeysWithEnv(t *testing.T) {
+	v := New()
+
+	// bind and define environment variables (including a nested one)
+	v.BindEnv("id")
+	v.BindEnv("foo.bar")
+	v.SetEnvKeyReplacer(strings.NewReplacer(".", "_"))
+	os.Setenv("ID", "13")
+	os.Setenv("FOO_BAR", "baz")
+
+	expectedKeys := sort.StringSlice{"id", "foo.bar"}
+	expectedKeys.Sort()
+	keys := sort.StringSlice(v.AllKeys())
+	keys.Sort()
+	assert.Equal(t, expectedKeys, keys)
+}
+
 func TestAliasesOfAliases(t *testing.T) {
 	Set("Title", "Checking Case")
 	RegisterAlias("Foo", "Bar")