[cmd] added a helper module cli for clean cmd (#3235)

* [cmd] added a helper module cli for clean cmd

* [cmd] added parse.go and register.go in internal/cli

* [cmd] fix golint issues

go.mod

@@ -57,6 +57,8 @@ require (
 	github.com/rs/cors v1.7.0 // indirect
 	github.com/rs/zerolog v1.18.0
 	github.com/shirou/gopsutil v2.18.12+incompatible // indirect
+	github.com/spf13/cobra v0.0.5
+	github.com/spf13/pflag v1.0.5
 	github.com/spf13/viper v1.6.1
 	github.com/stretchr/testify v1.6.1
 	github.com/syndtr/goleveldb v1.0.1-0.20190923125748-758128399b1d

internal/cli/flag.go

@@ -0,0 +1,93 @@
+// cli is the module for customized cli structures for cobra command.
+// The module is a wrapper over spf13/pflag, and serve to clean the code structure.
+
+package cli
+
+import "github.com/spf13/pflag"
+
+// Flag is the interface for cli flags.
+// To get the value after cli parsing, use fs.GetString(flag.Name)
+type Flag interface {
+	RegisterTo(fs *pflag.FlagSet) error
+}
+
+// StringFlag is the flag with string value
+type StringFlag struct {
+	Name       string
+	Shorthand  string
+	Usage      string
+	Deprecated string
+	Hidden     bool
+
+	DefValue string
+}
+
+// RegisterTo register the string flag to FlagSet
+func (f StringFlag) RegisterTo(fs *pflag.FlagSet) error {
+	fs.StringP(f.Name, f.Shorthand, f.DefValue, f.Usage)
+	return markHiddenOrDeprecated(fs, f.Name, f.Deprecated, f.Hidden)
+}
+
+// BoolFlag is the flag with boolean value
+type BoolFlag struct {
+	Name       string
+	Shorthand  string
+	Usage      string
+	Deprecated string
+	Hidden     bool
+
+	DefValue bool
+}
+
+// RegisterTo register the boolean flag to FlagSet
+func (f BoolFlag) RegisterTo(fs *pflag.FlagSet) error {
+	fs.BoolP(f.Name, f.Shorthand, f.DefValue, f.Usage)
+	return markHiddenOrDeprecated(fs, f.Name, f.Deprecated, f.Hidden)
+}
+
+// IntFlag is the flag with int value
+type IntFlag struct {
+	Name       string
+	Shorthand  string
+	Usage      string
+	Deprecated string
+	Hidden     bool
+
+	DefValue int
+}
+
+// RegisterTo register the int flag to FlagSet
+func (f IntFlag) RegisterTo(fs *pflag.FlagSet) error {
+	fs.IntP(f.Name, f.Shorthand, f.DefValue, f.Usage)
+	return markHiddenOrDeprecated(fs, f.Name, f.Deprecated, f.Hidden)
+}
+
+// StringSliceFlag is the flag with string slice value
+type StringSliceFlag struct {
+	Name       string
+	Shorthand  string
+	Usage      string
+	Deprecated string
+	Hidden     bool
+
+	DefValue []string
+}
+
+// RegisterTo register the string slice flag to FlagSet
+func (f StringSliceFlag) RegisterTo(fs *pflag.FlagSet) error {
+	fs.StringSliceP(f.Name, f.Shorthand, f.DefValue, f.Usage)
+	return markHiddenOrDeprecated(fs, f.Name, f.Deprecated, f.Hidden)
+}
+
+func markHiddenOrDeprecated(fs *pflag.FlagSet, name string, deprecated string, hidden bool) error {
+	if len(deprecated) != 0 {
+		if err := fs.MarkDeprecated(name, deprecated); err != nil {
+			return err
+		}
+	} else if hidden {
+		if err := fs.MarkHidden(name); err != nil {
+			return err
+		}
+	}
+	return nil
+}

internal/cli/flag_test.go

@@ -0,0 +1,74 @@
+package cli
+
+import (
+	"testing"
+
+	"github.com/spf13/cobra"
+)
+
+func TestStringFlag(t *testing.T) {
+	tests := []struct {
+		flagName   string
+		deprecated string
+		hidden     bool
+		defValue   string
+		args       []string
+		expValue   string
+	}{
+		{
+			defValue: "default",
+			args:     []string{},
+			expValue: "default",
+		},
+		{
+			defValue: "default",
+			args:     []string{"--test", "custom"},
+			expValue: "custom",
+		},
+		{
+			defValue:   "default",
+			args:       []string{},
+			deprecated: "deprecated",
+			expValue:   "default",
+		},
+		{
+			defValue: "default",
+			args:     []string{},
+			hidden:   true,
+			expValue: "default",
+		},
+	}
+	for i, test := range tests {
+		flagName := "test"
+		flag := StringFlag{
+			Name:       flagName,
+			Deprecated: test.deprecated,
+			Hidden:     test.hidden,
+			DefValue:   test.defValue,
+		}
+		var (
+			gotValue string
+			err      error
+		)
+		cmd := makeTestCommand(func(cmd *cobra.Command, args []string) {
+			gotValue, err = cmd.Flags().GetString(flagName)
+			if err != nil {
+				t.Fatalf("Test %v: %v", i, err)
+			}
+		})
+		if err := flag.RegisterTo(cmd.Flags()); err != nil {
+			t.Fatalf("Test %v: %v", i, err)
+		}
+		cmd.SetArgs(test.args)
+		if err := cmd.Execute(); err != nil {
+			t.Fatalf("Test %v: %v", i, err)
+		}
+		if gotValue != test.expValue {
+			t.Errorf("Test %v: unexpected flag value %v / %v", i, gotValue, test.expValue)
+		}
+	}
+}
+
+func makeTestCommand(run func(cmd *cobra.Command, args []string)) *cobra.Command {
+	return &cobra.Command{Run: run}
+}

internal/cli/parse.go

@@ -0,0 +1,112 @@
+package cli
+
+import (
+	"github.com/spf13/cobra"
+	"github.com/spf13/pflag"
+)
+
+type errorHandle func(error)
+
+var (
+	parseErrorHandleFunc errorHandle
+)
+
+// SetParseErrorHandle set the error handle function used for cli parsing flags.
+// An error handle example:
+// 		cli.SetParseErrorHandle(func(err error) {
+//			fmt.Println(err)
+// 			os.Exit(3)
+// 		})
+func SetParseErrorHandle(f errorHandle) {
+	parseErrorHandleFunc = f
+}
+
+// GetStringFlagValue get the string value for the given StringFlag from the local flags
+// of the cobra command.
+func GetStringFlagValue(cmd *cobra.Command, flag StringFlag) string {
+	return getStringFlagValue(cmd.Flags(), flag)
+}
+
+// GetStringPersistentFlagValue get the string value for the given StringFlag from the persistent
+// flags of the cobra command.
+func GetStringPersistentFlagValue(cmd *cobra.Command, flag StringFlag) string {
+	return getStringFlagValue(cmd.PersistentFlags(), flag)
+}
+
+func getStringFlagValue(fs *pflag.FlagSet, flag StringFlag) string {
+	val, err := fs.GetString(flag.Name)
+	if err != nil {
+		handleParseError(err)
+		return ""
+	}
+	return val
+}
+
+// GetBoolFlagValue get the bool value for the given BoolFlag from the local flags of the
+// cobra command.
+func GetBoolFlagValue(cmd *cobra.Command, flag BoolFlag) bool {
+	return getBoolFlagValue(cmd.Flags(), flag)
+}
+
+// GetBoolPersistentFlagValue get the bool value for the given BoolFlag from the persistent flags
+// of the given cobra command.
+func GetBoolPersistentFlagValue(cmd *cobra.Command, flag BoolFlag) bool {
+	return getBoolFlagValue(cmd.PersistentFlags(), flag)
+}
+
+func getBoolFlagValue(fs *pflag.FlagSet, flag BoolFlag) bool {
+	val, err := fs.GetBool(flag.Name)
+	if err != nil {
+		handleParseError(err)
+		return false
+	}
+	return val
+}
+
+// GetIntFlagValue get the int value for the given IntFlag from the local flags of the
+// cobra command.
+func GetIntFlagValue(cmd *cobra.Command, flag IntFlag) int {
+	return getIntFlagValue(cmd.Flags(), flag)
+}
+
+// GetIntPersistentFlagValue get the int value for the given IntFlag from the persistent
+// flags of the cobra command.
+func GetIntPersistentFlagValue(cmd *cobra.Command, flag IntFlag) int {
+	return getIntFlagValue(cmd.PersistentFlags(), flag)
+}
+
+func getIntFlagValue(fs *pflag.FlagSet, flag IntFlag) int {
+	val, err := fs.GetInt(flag.Name)
+	if err != nil {
+		handleParseError(err)
+		return 0
+	}
+	return val
+}
+
+// GetStringSliceFlagValue get the string slice value for the given StringSliceFlag from
+// the local flags of the cobra command.
+func GetStringSliceFlagValue(cmd *cobra.Command, flag StringSliceFlag) []string {
+	return getStringSliceFlagValue(cmd.Flags(), flag)
+}
+
+// GetStringSlicePersistentFlagValue get the string slice value for the given StringSliceFlag
+// from the persistent flags of the cobra command.
+func GetStringSlicePersistentFlagValue(cmd *cobra.Command, flag StringSliceFlag) []string {
+	return getStringSliceFlagValue(cmd.PersistentFlags(), flag)
+}
+
+func getStringSliceFlagValue(fs *pflag.FlagSet, flag StringSliceFlag) []string {
+	val, err := fs.GetStringSlice(flag.Name)
+	if err != nil {
+		handleParseError(err)
+		return nil
+	}
+	return val
+}
+
+func handleParseError(err error) {
+	if parseErrorHandleFunc != nil {
+		parseErrorHandleFunc(err)
+	}
+}

internal/cli/parse_test.go

@@ -0,0 +1,152 @@
+package cli
+
+import (
+	"bytes"
+	"errors"
+	"fmt"
+	"strings"
+	"testing"
+
+	"github.com/spf13/cobra"
+)
+
+var (
+	testStringFlag1 = StringFlag{
+		Name:     "string1",
+		DefValue: "default",
+	}
+
+	testStringFlag2 = StringFlag{
+		Name:     "string2",
+		DefValue: "default",
+	}
+
+	testIntFlag1 = IntFlag{
+		Name:     "int1",
+		DefValue: 12345,
+	}
+
+	testIntFlag2 = IntFlag{
+		Name:     "int2",
+		DefValue: 12345,
+	}
+
+	testBoolFlag1 = BoolFlag{
+		Name:     "bool1",
+		DefValue: true,
+	}
+
+	testBoolFlag2 = BoolFlag{
+		Name:     "bool2",
+		DefValue: true,
+	}
+
+	testStringSliceFlag1 = StringSliceFlag{
+		Name:     "string-slice1",
+		DefValue: []string{"string1", "string2"},
+	}
+
+	testStringSliceFlag2 = StringSliceFlag{
+		Name:     "string-slice2",
+		DefValue: []string{"string1", "string2"},
+	}
+)
+
+func TestGetStringFlagValue(t *testing.T) {
+	defer tearDown()
+
+	tests := []struct {
+		flags     []Flag
+		pflags    []Flag
+		args      []string
+		queryFlag StringFlag
+
+		expExecErr  error
+		expParseErr error
+		expVal      string
+	}{
+		{
+			flags:     []Flag{testStringFlag1},
+			args:      []string{},
+			queryFlag: testStringFlag1,
+			expVal:    "default",
+		},
+		{
+			flags:     []Flag{testStringFlag1},
+			args:      []string{"--string1", "custom"},
+			queryFlag: testStringFlag1,
+			expVal:    "custom",
+		},
+		{
+			pflags:    []Flag{testStringFlag1},
+			args:      []string{"--string1", "custom"},
+			queryFlag: testStringFlag1,
+			expVal:    "custom",
+		},
+		{
+			flags:      []Flag{testStringFlag1},
+			args:       []string{"--string1"},
+			expExecErr: errors.New("flag needs an argument"),
+		},
+		{
+			flags:       []Flag{testStringFlag1},
+			args:        []string{},
+			queryFlag:   testStringFlag2,
+			expParseErr: errors.New("flag accessed but not defined"),
+		},
+	}
+	for i, test := range tests {
+		var (
+			val      string
+			parseErr error
+		)
+		SetParseErrorHandle(func(gotErr error) {
+			parseErr = gotErr
+		})
+		cmd := makeTestCommand(func(cmd *cobra.Command, args []string) {
+			val = GetStringFlagValue(cmd, test.queryFlag)
+		})
+		cmd.SetOut(bytes.NewBuffer(nil))
+		if err := RegisterFlags(cmd, test.flags); err != nil {
+			t.Fatal(err)
+		}
+		if err := RegisterPFlags(cmd, test.pflags); err != nil {
+			t.Fatal(err)
+		}
+		cmd.SetArgs(test.args)
+		execErr := cmd.Execute()
+
+		if err := assertError(execErr, test.expExecErr); err != nil {
+			t.Fatalf("Test %v execution: %v", i, err)
+		}
+		if execErr != nil || test.expExecErr != nil {
+			continue
+		}
+		if err := assertError(parseErr, test.expParseErr); err != nil {
+			t.Fatalf("Test %v parse: %v", i, err)
+		}
+		if parseErr != nil || test.expParseErr != nil {
+			continue
+		}
+		if val != test.expVal {
+			t.Errorf("Test %v: unexpected value %v / %v", i, val, test.expVal)
+		}
+	}
+}
+
+func tearDown() {
+	parseErrorHandleFunc = nil
+}
+
+func assertError(gotErr, expErr error) error {
+	if (gotErr == nil) != (expErr == nil) {
+		return fmt.Errorf("error unexpected [%v] / [%v]", gotErr, expErr)
+	}
+	if gotErr == nil {
+		return nil
+	}
+	if !strings.Contains(gotErr.Error(), expErr.Error()) {
+		return fmt.Errorf("error unexpected [%v] / [%v]", gotErr, expErr)
+	}
+	return nil
+}

internal/cli/register.go

@@ -0,0 +1,29 @@
+package cli
+
+import (
+	"github.com/spf13/cobra"
+)
+
+// RegisterFlags register the flags to command's local flag
+func RegisterFlags(cmd *cobra.Command, flags []Flag) error {
+	fs := cmd.Flags()
+
+	for _, flag := range flags {
+		if err := flag.RegisterTo(fs); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// RegisterPFlags register the flags to command's persistent flag
+func RegisterPFlags(cmd *cobra.Command, flags []Flag) error {
+	fs := cmd.PersistentFlags()
+
+	for _, flag := range flags {
+		if err := flag.RegisterTo(fs); err != nil {
+			return err
+		}
+	}
+	return nil
+}