[go] 環境変数からフラグをセットする

main.go

package main

import (
	"flag"
	"fmt"
	"os"
	"strings"
)

const (
	flag_prefix = "MYKEY"
)

func main() {
	var age int
	var fooBar string
	flag.IntVar(&age, "age", 0, "age")
	flag.StringVar(&fooBar, "foo-bar", "foobar", "foobar")

	// 環境変数から値を設定する。
	// コマンドライン引数を優先させるため、`flag.Parse`よりも前に呼び出す
	SetFlagFromEnvironment(flag.CommandLine, flag_prefix)

	flag.Parse()

	fmt.Printf("Age = %d\n", age)
	fmt.Printf("FooBar = %s\n", fooBar)
}

func SetFlagFromEnvironment(fs *flag.FlagSet, prefix string) {
	fs.VisitAll(func(f *flag.Flag) {
		envKey := strings.ToUpper(strings.Replace(f.Name, "-", "_", -1))
		if prefix != "" {
			envKey = prefix + "_" + envKey
		}

		if s := os.Getenv(envKey); s != "" {
			f.Value.Set(s)
		}
	})
}

main_test.go

package main

import (
	"flag"
	"os"
	"strings"
	"testing"
)

func TestSetFlagFromEnvironment(t *testing.T) {
	var envPrefix = "MYKEY"

	cases := []struct {
		name         string
		defaultValue string
		envValue     string
		expected     string
	}{
		{"foo-bar", "foobar", "hoge", "hoge"},
		{"age", "0", "", "0"},
	}

	for _, c := range cases {
		var v string
		fs := flag.NewFlagSet("", flag.ExitOnError)
		fs.StringVar(&v, c.name, c.defaultValue, "")

		envKey := envPrefix + "_" + strings.ToUpper(strings.Replace(c.name, "-", "_", -1))
		os.Setenv(envKey, c.envValue)

		SetFlagFromEnvironment(fs, envPrefix)

		if v != c.expected {
			t.Errorf("got %q, expected %q", v, c.expected)
		}
	}

}