diff --git a/cmd/influx/main.go b/cmd/influx/main.go index 9b614d0c592..2093e8bfa0d 100644 --- a/cmd/influx/main.go +++ b/cmd/influx/main.go @@ -55,15 +55,17 @@ func newHTTPClient() (*httpc.Client, error) { type ( cobraRunEFn func(cmd *cobra.Command, args []string) error - cobraRuneEMiddleware func(fn cobraRunEFn) cobraRunEFn + cobraRunEMiddleware func(fn cobraRunEFn) cobraRunEFn genericCLIOptFn func(*genericCLIOpts) ) type genericCLIOpts struct { - in io.Reader - w io.Writer - runEWrapFn cobraRuneEMiddleware + in io.Reader + w io.Writer + errW io.Writer + + runEWrapFn cobraRunEMiddleware } func (o genericCLIOpts) newCmd(use string, runE func(*cobra.Command, []string) error) *cobra.Command { @@ -76,6 +78,8 @@ func (o genericCLIOpts) newCmd(use string, runE func(*cobra.Command, []string) e cmd.RunE = o.runEWrapFn(runE) } cmd.SetOut(o.w) + cmd.SetIn(o.in) + cmd.SetErr(o.errW) return cmd } @@ -95,6 +99,18 @@ func out(w io.Writer) genericCLIOptFn { } } +func err(w io.Writer) genericCLIOptFn { + return func(o *genericCLIOpts) { + o.errW = w + } +} + +func runEMiddlware(mw cobraRunEMiddleware) genericCLIOptFn { + return func(o *genericCLIOpts) { + o.runEWrapFn = mw + } +} + type globalFlags struct { config.Config local bool @@ -115,6 +131,7 @@ func newInfluxCmdBuilder(optFns ...genericCLIOptFn) *cmdInfluxBuilder { opt := genericCLIOpts{ in: os.Stdin, w: os.Stdout, + errW: os.Stderr, runEWrapFn: checkSetupRunEMiddleware(&flags), } for _, optFn := range optFns { @@ -298,7 +315,7 @@ func checkSetup(host string, skipVerify bool) error { return nil } -func checkSetupRunEMiddleware(f *globalFlags) cobraRuneEMiddleware { +func checkSetupRunEMiddleware(f *globalFlags) cobraRunEMiddleware { return func(fn cobraRunEFn) cobraRunEFn { return func(cmd *cobra.Command, args []string) error { err := fn(cmd, args) diff --git a/cmd/influx/main_test.go b/cmd/influx/main_test.go new file mode 100644 index 00000000000..8972145c474 --- /dev/null +++ b/cmd/influx/main_test.go @@ -0,0 +1,94 @@ +package main + +import ( + "bytes" + "io/ioutil" + "testing" + + "github.com/influxdata/influxdb/cmd/influx/config" + "github.com/spf13/cobra" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func Test_influx_cmd(t *testing.T) { + tests := []struct { + name string + args []string + envVars map[string]string + expected globalFlags + }{ + { + name: "all full length flags set", + args: []string{"--token=TOKEN", "--host=HOST", "--local=true", "--skip-verify=true"}, + envVars: envVarsZeroMap, + expected: globalFlags{ + Config: config.Config{ + Token: "TOKEN", + Host: "HOST", + }, + skipVerify: true, + local: true, + }, + }, + { + name: "token p flag set", + args: []string{"-t=TOKEN", "--host=HOST", "--local=true", "--skip-verify=true"}, + envVars: envVarsZeroMap, + expected: globalFlags{ + Config: config.Config{ + Token: "TOKEN", + Host: "HOST", + }, + skipVerify: true, + local: true, + }, + }, + { + name: "env vars set", + args: []string{"--local=true", "--skip-verify=true"}, + envVars: map[string]string{ + "INFLUX_TOKEN": "TOKEN", + "INFLUX_HOST": "HOST", + }, + expected: globalFlags{ + Config: config.Config{ + Token: "TOKEN", + Host: "HOST", + }, + skipVerify: true, + local: true, + }, + }, + } + + for _, tt := range tests { + fn := func(t *testing.T) { + defer addEnvVars(t, tt.envVars)() + + builder := newInfluxCmdBuilder( + in(new(bytes.Buffer)), + out(ioutil.Discard), + err(ioutil.Discard), + runEMiddlware(func(fn cobraRunEFn) cobraRunEFn { return fn }), + ) + + flagCapture := new(globalFlags) + influxCmd := builder.cmd(func(f *globalFlags, opt genericCLIOpts) *cobra.Command { + flagCapture = f + return &cobra.Command{Use: "foo"} + }) + + influxCmd.SetArgs(append([]string{"foo"}, tt.args...)) + + require.NoError(t, influxCmd.Execute()) + + assert.Equal(t, tt.expected.Host, flagCapture.Host) + assert.Equal(t, tt.expected.Token, flagCapture.Token) + assert.Equal(t, tt.expected.local, flagCapture.local) + assert.Equal(t, tt.expected.skipVerify, flagCapture.skipVerify) + } + + t.Run(tt.name, fn) + } +}