feat: make Unmarshal work with AutomaticEnv

Co-authored-by: Filip Krakowski <krakowski@hhu.de>
Signed-off-by: Mark Sagi-Kazar <mark.sagikazar@gmail.com>
This commit is contained in:
Mark Sagi-Kazar 2023-12-05 15:40:27 +01:00 committed by Márk Sági-Kazár
parent 6ea31ae4ca
commit 73dfb94c57
2 changed files with 125 additions and 1 deletions

View file

@ -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

View file

@ -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)