From 7fc3ceecab7e0edcbdc4520f6437d1f9bdfecaec Mon Sep 17 00:00:00 2001
From: Haytham AbuelFutuh <habuelfutuh@lyft.com>
Date: Wed, 8 May 2019 13:39:26 -0700
Subject: [PATCH] InitialzePFlags doesn't add flags for sub-sections (#17)

* Allow InitializePFlags to recursively discover sub-section flags

* Add unit test
---
 config/tests/accessor_test.go | 30 ++++++++++++++++++++++++++++++
 config/viper/viper.go         | 13 +++++++++++--
 2 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/config/tests/accessor_test.go b/config/tests/accessor_test.go
index e6f9309..18a81e5 100644
--- a/config/tests/accessor_test.go
+++ b/config/tests/accessor_test.go
@@ -159,6 +159,36 @@ func TestAccessor_InitializePflags(t *testing.T) {
 			assert.Equal(t, 4, otherC.IntValue)
 			assert.Equal(t, []string{"default value"}, otherC.StringArrayWithDefaults)
 		})
+
+		t.Run(fmt.Sprintf("[%v] Sub-sections", provider(config.Options{}).ID()), func(t *testing.T) {
+			reg := config.NewRootSection()
+			sec, err := reg.RegisterSection(MyComponentSectionKey, &MyComponentConfig{})
+			assert.NoError(t, err)
+
+			_, err = sec.RegisterSection("nested", &OtherComponentConfig{})
+			assert.NoError(t, err)
+
+			v := provider(config.Options{
+				SearchPaths: []string{filepath.Join("testdata", "nested_config.yaml")},
+				RootSection: reg,
+			})
+
+			set := pflag.NewFlagSet("test", pflag.ExitOnError)
+			v.InitializePflags(set)
+			assert.NoError(t, set.Parse([]string{"--my-component.nested.int-val=3"}))
+			assert.True(t, set.Parsed())
+
+			flagValue, err := set.GetInt("my-component.nested.int-val")
+			assert.NoError(t, err)
+			assert.Equal(t, 3, flagValue)
+
+			assert.NoError(t, v.UpdateConfig(context.TODO()))
+			r := reg.GetSection(MyComponentSectionKey).GetConfig().(*MyComponentConfig)
+			assert.Equal(t, "Hello World", r.StringValue)
+
+			nested := sec.GetSection("nested").GetConfig().(*OtherComponentConfig)
+			assert.Equal(t, 3, nested.IntValue)
+		})
 	}
 }
 
diff --git a/config/viper/viper.go b/config/viper/viper.go
index 05a5378..c9df21f 100644
--- a/config/viper/viper.go
+++ b/config/viper/viper.go
@@ -64,9 +64,18 @@ func (v viperAccessor) InitializePflags(cmdFlags *pflag.FlagSet) {
 }
 
 func (v viperAccessor) addSectionsPFlags(flags *pflag.FlagSet) (err error) {
-	for key, section := range v.rootConfig.GetSections() {
+	return v.addSubsectionsPFlags(flags, "", v.rootConfig)
+}
+
+func (v viperAccessor) addSubsectionsPFlags(flags *pflag.FlagSet, rootKey string, root config.Section) error {
+	for key, section := range root.GetSections() {
+		prefix := rootKey + key + keyDelim
 		if asPFlagProvider, ok := section.GetConfig().(config.PFlagProvider); ok {
-			flags.AddFlagSet(asPFlagProvider.GetPFlagSet(key + keyDelim))
+			flags.AddFlagSet(asPFlagProvider.GetPFlagSet(prefix))
+		}
+
+		if err := v.addSubsectionsPFlags(flags, prefix, section); err != nil {
+			return err
 		}
 	}