pkg: move SetFlagsFromEnv to pkg package

pkg/flag.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"log"
 	"os"
+	"strings"
 )
 
 type DeprecatedFlag struct {
@@ -42,3 +43,24 @@ func UsageWithIgnoredFlagsFunc(fs *flag.FlagSet, ignore []string) func() {
 		})
 	}
 }
+
+// SetFlagsFromEnv parses all registered flags in the given flagset,
+// and if they are not already set it attempts to set their values from
+// environment variables. Environment variables take the name of the flag but
+// are UPPERCASE, have the prefix "ETCD_", and any dashes are replaced by
+// underscores - for example: some-flag => ETCD_SOME_FLAG
+func SetFlagsFromEnv(fs *flag.FlagSet) {
+	alreadySet := make(map[string]bool)
+	fs.Visit(func(f *flag.Flag) {
+		alreadySet[f.Name] = true
+	})
+	fs.VisitAll(func(f *flag.Flag) {
+		if !alreadySet[f.Name] {
+			key := "ETCD_" + strings.ToUpper(strings.Replace(f.Name, "-", "_", -1))
+			val := os.Getenv(key)
+			if val != "" {
+				fs.Set(f.Name, val)
+			}
+		}
+	})
+}

pkg/flag_test.go

@@ -0,0 +1,51 @@
+package pkg
+
+import (
+	"flag"
+	"os"
+	"testing"
+)
+
+func TestSetFlagsFromEnv(t *testing.T) {
+	fs := flag.NewFlagSet("testing", flag.ExitOnError)
+	fs.String("a", "", "")
+	fs.String("b", "", "")
+	fs.String("c", "", "")
+	fs.Parse([]string{})
+
+	os.Clearenv()
+	// flags should be settable using env vars
+	os.Setenv("ETCD_A", "foo")
+	// and command-line flags
+	if err := fs.Set("b", "bar"); err != nil {
+		t.Fatal(err)
+	}
+	// command-line flags take precedence over env vars
+	os.Setenv("ETCD_C", "woof")
+	if err := fs.Set("c", "quack"); err != nil {
+		t.Fatal(err)
+	}
+
+	// first verify that flags are as expected before reading the env
+	for f, want := range map[string]string{
+		"a": "",
+		"b": "bar",
+		"c": "quack",
+	} {
+		if got := fs.Lookup(f).Value.String(); got != want {
+			t.Fatalf("flag %q=%q, want %q", f, got, want)
+		}
+	}
+
+	// now read the env and verify flags were updated as expected
+	SetFlagsFromEnv(fs)
+	for f, want := range map[string]string{
+		"a": "foo",
+		"b": "bar",
+		"c": "quack",
+	} {
+		if got := fs.Lookup(f).Value.String(); got != want {
+			t.Errorf("flag %q=%q, want %q", f, got, want)
+		}
+	}
+}