allow targeted TF_CLI_ARGS_x

This commit is contained in:
Mitchell Hashimoto 2017-02-13 15:12:29 -08:00
parent 53796fcdb0
commit df93e5120c
No known key found for this signature in database
GPG Key ID: 744E147AA52F5B0A
2 changed files with 142 additions and 33 deletions

93
main.go
View File

@ -7,6 +7,7 @@ import (
"log" "log"
"os" "os"
"runtime" "runtime"
"strings"
"sync" "sync"
"github.com/hashicorp/go-plugin" "github.com/hashicorp/go-plugin"
@ -138,39 +139,27 @@ func wrappedMain() int {
// Get the command line args. // Get the command line args.
args := os.Args[1:] args := os.Args[1:]
// Build the CLI so far, we do this so we can query the subcommand.
cliRunner := &cli.CLI{
Args: args,
Commands: Commands,
HelpFunc: helpFunc,
HelpWriter: os.Stdout,
}
// Prefix the args with any args from the EnvCLI // Prefix the args with any args from the EnvCLI
if v := os.Getenv(EnvCLI); v != "" { args, err = mergeEnvArgs(EnvCLI, args)
log.Printf("[INFO] %s value: %q", EnvCLI, v) if err != nil {
extra, err := shellwords.Parse(v) Ui.Error(err.Error())
if err != nil { return 1
Ui.Error(fmt.Sprintf( }
"Error parsing extra CLI args from %s: %s",
EnvCLI, err))
return 1
}
// Find the index to place the flags. We put them exactly
// after the first non-flag arg.
idx := -1
for i, v := range args {
if len(v) > 0 && v[0] != '-' {
idx = i
break
}
}
// idx points to the exact arg that isn't a flag. We increment
// by one so that all the copying below expects idx to be the
// insertion point.
idx++
// Copy the args
newArgs := make([]string, len(args)+len(extra))
copy(newArgs, args[:idx])
copy(newArgs[idx:], extra)
copy(newArgs[len(extra)+idx:], args[idx:])
args = newArgs
// Prefix the args with any args from the EnvCLI targeting this command
suffix := strings.Replace(cliRunner.Subcommand(), "-", "_", -1)
args, err = mergeEnvArgs(fmt.Sprintf("%s_%s", EnvCLI, suffix), args)
if err != nil {
Ui.Error(err.Error())
return 1
} }
// We shortcut "--version" and "-v" to just show the version // We shortcut "--version" and "-v" to just show the version
@ -184,8 +173,9 @@ func wrappedMain() int {
} }
} }
// Rebuild the CLI with any modified args.
log.Printf("[INFO] CLI command args: %#v", args) log.Printf("[INFO] CLI command args: %#v", args)
cli := &cli.CLI{ cliRunner = &cli.CLI{
Args: args, Args: args,
Commands: Commands, Commands: Commands,
HelpFunc: helpFunc, HelpFunc: helpFunc,
@ -196,7 +186,7 @@ func wrappedMain() int {
ContextOpts.Providers = config.ProviderFactories() ContextOpts.Providers = config.ProviderFactories()
ContextOpts.Provisioners = config.ProvisionerFactories() ContextOpts.Provisioners = config.ProvisionerFactories()
exitCode, err := cli.Run() exitCode, err := cliRunner.Run()
if err != nil { if err != nil {
Ui.Error(fmt.Sprintf("Error executing CLI: %s", err.Error())) Ui.Error(fmt.Sprintf("Error executing CLI: %s", err.Error()))
return 1 return 1
@ -284,3 +274,40 @@ func copyOutput(r io.Reader, doneCh chan<- struct{}) {
wg.Wait() wg.Wait()
} }
func mergeEnvArgs(envName string, args []string) ([]string, error) {
v := os.Getenv(envName)
if v == "" {
return args, nil
}
log.Printf("[INFO] %s value: %q", envName, v)
extra, err := shellwords.Parse(v)
if err != nil {
return nil, fmt.Errorf(
"Error parsing extra CLI args from %s: %s",
envName, err)
}
// Find the index to place the flags. We put them exactly
// after the first non-flag arg.
idx := -1
for i, v := range args {
if len(v) > 0 && v[0] != '-' {
idx = i
break
}
}
// idx points to the exact arg that isn't a flag. We increment
// by one so that all the copying below expects idx to be the
// insertion point.
idx++
// Copy the args
newArgs := make([]string, len(args)+len(extra))
copy(newArgs, args[:idx])
copy(newArgs[idx:], extra)
copy(newArgs[len(extra)+idx:], args[idx:])
return newArgs, nil
}

View File

@ -110,6 +110,7 @@ func TestMain_cliArgsFromEnv(t *testing.T) {
for i, tc := range cases { for i, tc := range cases {
t.Run(fmt.Sprintf("%d-%s", i, tc.Name), func(t *testing.T) { t.Run(fmt.Sprintf("%d-%s", i, tc.Name), func(t *testing.T) {
os.Unsetenv(EnvCLI) os.Unsetenv(EnvCLI)
defer os.Unsetenv(EnvCLI)
// Set the env var value // Set the env var value
if tc.Value != "" { if tc.Value != "" {
@ -142,6 +143,87 @@ func TestMain_cliArgsFromEnv(t *testing.T) {
} }
} }
// This test just has more options than the test above. Use this for
// more control over behavior at the expense of more complex test structures.
func TestMain_cliArgsFromEnvAdvanced(t *testing.T) {
// Restore original CLI args
oldArgs := os.Args
defer func() { os.Args = oldArgs }()
cases := []struct {
Name string
Command string
EnvVar string
Args []string
Value string
Expected []string
Err bool
}{
{
"targeted to another command",
"command",
EnvCLI + "_foo",
[]string{"command", "foo", "bar"},
"-flag",
[]string{"foo", "bar"},
false,
},
{
"targeted to this command",
"command",
EnvCLI + "_command",
[]string{"command", "foo", "bar"},
"-flag",
[]string{"-flag", "foo", "bar"},
false,
},
}
for i, tc := range cases {
t.Run(fmt.Sprintf("%d-%s", i, tc.Name), func(t *testing.T) {
// Setup test command and restore that
testCommandName := tc.Command
testCommand := &testCommandCLI{}
defer func() { delete(Commands, testCommandName) }()
Commands[testCommandName] = func() (cli.Command, error) {
return testCommand, nil
}
os.Unsetenv(tc.EnvVar)
defer os.Unsetenv(tc.EnvVar)
// Set the env var value
if tc.Value != "" {
if err := os.Setenv(tc.EnvVar, tc.Value); err != nil {
t.Fatalf("err: %s", err)
}
}
// Setup the args
args := make([]string, len(tc.Args)+1)
args[0] = oldArgs[0] // process name
copy(args[1:], tc.Args)
// Run it!
os.Args = args
testCommand.Args = nil
exit := wrappedMain()
if (exit != 0) != tc.Err {
t.Fatalf("bad: %d", exit)
}
if tc.Err {
return
}
// Verify
if !reflect.DeepEqual(testCommand.Args, tc.Expected) {
t.Fatalf("bad: %#v", testCommand.Args)
}
})
}
}
type testCommandCLI struct { type testCommandCLI struct {
Args []string Args []string
} }