chore: Add helper for uniform flags and env vars (#588)

2026-09-24 15:04:27 +08:00 · 2022-03-28 14:26:41 -05:00
parent be8389fd74
commit 01957da040
4 changed files with 248 additions and 57 deletions
@@ -0,0 +1,70 @@
+// Package cliflag extends flagset with environment variable defaults.
+//
+// Usage:
+//
+// cliflag.String(root.Flags(), &address, "address", "a", "CODER_ADDRESS", "127.0.0.1:3000", "The address to serve the API and dashboard")
+//
+// Will produce the following usage docs:
+//
+//   -a, --address string              The address to serve the API and dashboard (uses $CODER_ADDRESS). (default "127.0.0.1:3000")
+//
+package cliflag
+
+import (
+	"fmt"
+	"os"
+	"strconv"
+
+	"github.com/spf13/pflag"
+)
+
+// StringVarP sets a string flag on the given flag set.
+func StringVarP(flagset *pflag.FlagSet, p *string, name string, shorthand string, env string, def string, usage string) {
+	v, ok := os.LookupEnv(env)
+	if !ok || v == "" {
+		v = def
+	}
+	flagset.StringVarP(p, name, shorthand, v, fmtUsage(usage, env))
+}
+
+// Uint8VarP sets a uint8 flag on the given flag set.
+func Uint8VarP(flagset *pflag.FlagSet, ptr *uint8, name string, shorthand string, env string, def uint8, usage string) {
+	val, ok := os.LookupEnv(env)
+	if !ok || val == "" {
+		flagset.Uint8VarP(ptr, name, shorthand, def, fmtUsage(usage, env))
+		return
+	}
+
+	vi64, err := strconv.ParseUint(val, 10, 8)
+	if err != nil {
+		flagset.Uint8VarP(ptr, name, shorthand, def, fmtUsage(usage, env))
+		return
+	}
+
+	flagset.Uint8VarP(ptr, name, shorthand, uint8(vi64), fmtUsage(usage, env))
+}
+
+// BoolVarP sets a bool flag on the given flag set.
+func BoolVarP(flagset *pflag.FlagSet, ptr *bool, name string, shorthand string, env string, def bool, usage string) {
+	val, ok := os.LookupEnv(env)
+	if !ok || val == "" {
+		flagset.BoolVarP(ptr, name, shorthand, def, fmtUsage(usage, env))
+		return
+	}
+
+	valb, err := strconv.ParseBool(val)
+	if err != nil {
+		flagset.BoolVarP(ptr, name, shorthand, def, fmtUsage(usage, env))
+		return
+	}
+
+	flagset.BoolVarP(ptr, name, shorthand, valb, fmtUsage(usage, env))
+}
+
+func fmtUsage(u string, env string) string {
+	if env == "" {
+		return fmt.Sprintf("%s.", u)
+	}
+
+	return fmt.Sprintf("%s - consumes $%s.", u, env)
+}
@@ -0,0 +1,145 @@
+package cliflag_test
+
+import (
+	"fmt"
+	"strconv"
+	"testing"
+
+	"github.com/spf13/pflag"
+	"github.com/stretchr/testify/require"
+
+	"github.com/coder/coder/cli/cliflag"
+	"github.com/coder/coder/cryptorand"
+)
+
+// Testcliflag cannot run in parallel because it uses t.Setenv.
+//nolint:paralleltest
+func TestCliflag(t *testing.T) {
+	t.Run("StringDefault", func(t *testing.T) {
+		var p string
+		flagset, name, shorthand, env, usage := randomFlag()
+		def, _ := cryptorand.String(10)
+
+		cliflag.StringVarP(flagset, &p, name, shorthand, env, def, usage)
+		got, err := flagset.GetString(name)
+		require.NoError(t, err)
+		require.Equal(t, def, got)
+		require.Contains(t, flagset.FlagUsages(), usage)
+		require.Contains(t, flagset.FlagUsages(), fmt.Sprintf(" - consumes $%s", env))
+	})
+
+	t.Run("StringEnvVar", func(t *testing.T) {
+		var p string
+		flagset, name, shorthand, env, usage := randomFlag()
+		envValue, _ := cryptorand.String(10)
+		t.Setenv(env, envValue)
+		def, _ := cryptorand.String(10)
+
+		cliflag.StringVarP(flagset, &p, name, shorthand, env, def, usage)
+		got, err := flagset.GetString(name)
+		require.NoError(t, err)
+		require.Equal(t, envValue, got)
+	})
+
+	t.Run("EmptyEnvVar", func(t *testing.T) {
+		var p string
+		flagset, name, shorthand, _, usage := randomFlag()
+		def, _ := cryptorand.String(10)
+
+		cliflag.StringVarP(flagset, &p, name, shorthand, "", def, usage)
+		got, err := flagset.GetString(name)
+		require.NoError(t, err)
+		require.Equal(t, def, got)
+		require.Contains(t, flagset.FlagUsages(), usage)
+		require.NotContains(t, flagset.FlagUsages(), " - consumes")
+	})
+
+	t.Run("IntDefault", func(t *testing.T) {
+		var p uint8
+		flagset, name, shorthand, env, usage := randomFlag()
+		def, _ := cryptorand.Int63n(10)
+
+		cliflag.Uint8VarP(flagset, &p, name, shorthand, env, uint8(def), usage)
+		got, err := flagset.GetUint8(name)
+		require.NoError(t, err)
+		require.Equal(t, uint8(def), got)
+		require.Contains(t, flagset.FlagUsages(), usage)
+		require.Contains(t, flagset.FlagUsages(), fmt.Sprintf(" - consumes $%s", env))
+	})
+
+	t.Run("IntEnvVar", func(t *testing.T) {
+		var p uint8
+		flagset, name, shorthand, env, usage := randomFlag()
+		envValue, _ := cryptorand.Int63n(10)
+		t.Setenv(env, strconv.FormatUint(uint64(envValue), 10))
+		def, _ := cryptorand.Int()
+
+		cliflag.Uint8VarP(flagset, &p, name, shorthand, env, uint8(def), usage)
+		got, err := flagset.GetUint8(name)
+		require.NoError(t, err)
+		require.Equal(t, uint8(envValue), got)
+	})
+
+	t.Run("IntFailParse", func(t *testing.T) {
+		var p uint8
+		flagset, name, shorthand, env, usage := randomFlag()
+		envValue, _ := cryptorand.String(10)
+		t.Setenv(env, envValue)
+		def, _ := cryptorand.Int63n(10)
+
+		cliflag.Uint8VarP(flagset, &p, name, shorthand, env, uint8(def), usage)
+		got, err := flagset.GetUint8(name)
+		require.NoError(t, err)
+		require.Equal(t, uint8(def), got)
+	})
+
+	t.Run("BoolDefault", func(t *testing.T) {
+		var p bool
+		flagset, name, shorthand, env, usage := randomFlag()
+		def, _ := cryptorand.Bool()
+
+		cliflag.BoolVarP(flagset, &p, name, shorthand, env, def, usage)
+		got, err := flagset.GetBool(name)
+		require.NoError(t, err)
+		require.Equal(t, def, got)
+		require.Contains(t, flagset.FlagUsages(), usage)
+		require.Contains(t, flagset.FlagUsages(), fmt.Sprintf(" - consumes $%s", env))
+	})
+
+	t.Run("BoolEnvVar", func(t *testing.T) {
+		var p bool
+		flagset, name, shorthand, env, usage := randomFlag()
+		envValue, _ := cryptorand.Bool()
+		t.Setenv(env, strconv.FormatBool(envValue))
+		def, _ := cryptorand.Bool()
+
+		cliflag.BoolVarP(flagset, &p, name, shorthand, env, def, usage)
+		got, err := flagset.GetBool(name)
+		require.NoError(t, err)
+		require.Equal(t, envValue, got)
+	})
+
+	t.Run("BoolFailParse", func(t *testing.T) {
+		var p bool
+		flagset, name, shorthand, env, usage := randomFlag()
+		envValue, _ := cryptorand.String(10)
+		t.Setenv(env, envValue)
+		def, _ := cryptorand.Bool()
+
+		cliflag.BoolVarP(flagset, &p, name, shorthand, env, def, usage)
+		got, err := flagset.GetBool(name)
+		require.NoError(t, err)
+		require.Equal(t, def, got)
+	})
+}
+
+func randomFlag() (*pflag.FlagSet, string, string, string, string) {
+	fsname, _ := cryptorand.String(10)
+	flagset := pflag.NewFlagSet(fsname, pflag.PanicOnError)
+	name, _ := cryptorand.String(10)
+	shorthand, _ := cryptorand.String(1)
+	env, _ := cryptorand.String(10)
+	usage, _ := cryptorand.String(10)
+
+	return flagset, name, shorthand, env, usage
+}