From 73dfb94c57ad48bcdf3d40b7f47d69c2962d800d Mon Sep 17 00:00:00 2001
From: Mark Sagi-Kazar <mark.sagikazar@gmail.com>
Date: Tue, 5 Dec 2023 15:40:27 +0100
Subject: [PATCH] feat: make Unmarshal work with AutomaticEnv

Co-authored-by: Filip Krakowski <krakowski@hhu.de>
Signed-off-by: Mark Sagi-Kazar <mark.sagikazar@gmail.com>
---
 viper.go      | 27 +++++++++++++-
 viper_test.go | 99 +++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 125 insertions(+), 1 deletion(-)

diff --git a/viper.go b/viper.go
index 5cbcc9f..53ded8a 100644
--- a/viper.go
+++ b/viper.go
@@ -1111,7 +1111,32 @@ func Unmarshal(rawVal any, opts ...DecoderConfigOption) error {
 }
 
 func (v *Viper) Unmarshal(rawVal any, opts ...DecoderConfigOption) error {
-	return decode(v.AllSettings(), defaultDecoderConfig(rawVal, opts...))
+	// TODO: make this optional?
+	structKeys, err := v.decodeStructKeys(rawVal, opts...)
+	if err != nil {
+		return err
+	}
+
+	// TODO: struct keys should be enough?
+	return decode(v.getSettings(append(v.AllKeys(), structKeys...)), defaultDecoderConfig(rawVal, opts...))
+}
+
+func (v *Viper) decodeStructKeys(input any, opts ...DecoderConfigOption) ([]string, error) {
+	var structKeyMap map[string]any
+
+	err := decode(input, defaultDecoderConfig(&structKeyMap, opts...))
+	if err != nil {
+		return nil, err
+	}
+
+	flattenedStructKeyMap := v.flattenAndMergeMap(map[string]bool{}, structKeyMap, "")
+
+	r := make([]string, 0, len(flattenedStructKeyMap))
+	for v := range flattenedStructKeyMap {
+		r = append(r, v)
+	}
+
+	return r, nil
 }
 
 // defaultDecoderConfig returns default mapstructure.DecoderConfig with support
diff --git a/viper_test.go b/viper_test.go
index 0e416e7..b8274a9 100644
--- a/viper_test.go
+++ b/viper_test.go
@@ -948,6 +948,105 @@ func TestUnmarshalWithDecoderOptions(t *testing.T) {
 	}, &C)
 }
 
+func TestUnmarshalWithAutomaticEnv(t *testing.T) {
+	t.Setenv("PORT", "1313")
+	t.Setenv("NAME", "Steve")
+	t.Setenv("DURATION", "1s1ms")
+	t.Setenv("MODES", "1,2,3")
+	t.Setenv("SECRET", "42")
+	t.Setenv("FILESYSTEM_SIZE", "4096")
+
+	type AuthConfig struct {
+		Secret string `mapstructure:"secret"`
+	}
+
+	type StorageConfig struct {
+		Size int `mapstructure:"size"`
+	}
+
+	type Configuration struct {
+		Port     int           `mapstructure:"port"`
+		Name     string        `mapstructure:"name"`
+		Duration time.Duration `mapstructure:"duration"`
+
+		// Infer name from struct
+		Modes []int
+
+		// Squash nested struct (omit prefix)
+		Authentication AuthConfig `mapstructure:",squash"`
+
+		// Different key
+		Storage StorageConfig `mapstructure:"filesystem"`
+
+		// Omitted field
+		Flag bool `mapstructure:"flag"`
+	}
+
+	v := New()
+	v.SetEnvKeyReplacer(strings.NewReplacer(".", "_"))
+	v.AutomaticEnv()
+
+	t.Run("OK", func(t *testing.T) {
+		var config Configuration
+		if err := v.Unmarshal(&config); err != nil {
+			t.Fatalf("unable to decode into struct, %v", err)
+		}
+
+		assert.Equal(
+			t,
+			Configuration{
+				Name:     "Steve",
+				Port:     1313,
+				Duration: time.Second + time.Millisecond,
+				Modes:    []int{1, 2, 3},
+				Authentication: AuthConfig{
+					Secret: "42",
+				},
+				Storage: StorageConfig{
+					Size: 4096,
+				},
+			},
+			config,
+		)
+	})
+
+	t.Run("Precedence", func(t *testing.T) {
+		var config Configuration
+
+		v.Set("port", 1234)
+		if err := v.Unmarshal(&config); err != nil {
+			t.Fatalf("unable to decode into struct, %v", err)
+		}
+
+		assert.Equal(
+			t,
+			Configuration{
+				Name:     "Steve",
+				Port:     1234,
+				Duration: time.Second + time.Millisecond,
+				Modes:    []int{1, 2, 3},
+				Authentication: AuthConfig{
+					Secret: "42",
+				},
+				Storage: StorageConfig{
+					Size: 4096,
+				},
+			},
+			config,
+		)
+	})
+
+	t.Run("Unset", func(t *testing.T) {
+		var config Configuration
+
+		err := v.Unmarshal(&config, func(config *mapstructure.DecoderConfig) {
+			config.ErrorUnset = true
+		})
+
+		assert.Error(t, err, "expected viper.Unmarshal to return error due to unset field 'FLAG'")
+	})
+}
+
 func TestBindPFlags(t *testing.T) {
 	v := New() // create independent Viper object
 	flagSet := pflag.NewFlagSet("test", pflag.ContinueOnError)