main: add TF_CLI_ARGS to specify additional CLI args

This commit is contained in:
Mitchell Hashimoto 2017-02-13 14:05:37 -08:00
parent c36b6c42ba
commit 7f67b32169
No known key found for this signature in database
GPG Key ID: 744E147AA52F5B0A
2 changed files with 184 additions and 2 deletions

47
main.go
View File

@ -13,11 +13,17 @@ import (
"github.com/hashicorp/terraform/helper/logging"
"github.com/hashicorp/terraform/terraform"
"github.com/mattn/go-colorable"
"github.com/mattn/go-shellwords"
"github.com/mitchellh/cli"
"github.com/mitchellh/panicwrap"
"github.com/mitchellh/prefixedio"
)
const (
// EnvCLI is the environment variable name to set additional CLI args.
EnvCLI = "TF_CLI_ARGS"
)
func main() {
// Override global prefix set by go-dynect during init()
log.SetPrefix("")
@ -129,9 +135,45 @@ func wrappedMain() int {
// Make sure we clean up any managed plugins at the end of this
defer plugin.CleanupClients()
// Get the command line args. We shortcut "--version" and "-v" to
// just show the version.
// Get the command line args.
args := os.Args[1:]
// Prefix the args with any args from the EnvCLI
if v := os.Getenv(EnvCLI); v != "" {
log.Printf("[INFO] %s value: %q", EnvCLI, v)
extra, err := shellwords.Parse(v)
if err != nil {
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 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
}
// We shortcut "--version" and "-v" to just show the version
for _, arg := range args {
if arg == "-v" || arg == "-version" || arg == "--version" {
newArgs := make([]string, len(args)+1)
@ -142,6 +184,7 @@ func wrappedMain() int {
}
}
log.Printf("[INFO] CLI command args: %#v", args)
cli := &cli.CLI{
Args: args,
Commands: Commands,

139
main_test.go Normal file
View File

@ -0,0 +1,139 @@
package main
import (
"fmt"
"os"
"reflect"
"testing"
"github.com/mitchellh/cli"
)
func TestMain_cliArgsFromEnv(t *testing.T) {
// Setup the state. This test really messes with the environment and
// global state so we set things up to be restored.
// Restore original CLI args
oldArgs := os.Args
defer func() { os.Args = oldArgs }()
// Setup test command and restore that
testCommandName := "unit-test-cli-args"
testCommand := &testCommandCLI{}
defer func() { delete(Commands, testCommandName) }()
Commands[testCommandName] = func() (cli.Command, error) {
return testCommand, nil
}
cases := []struct {
Name string
Args []string
Value string
Expected []string
Err bool
}{
{
"no env",
[]string{testCommandName, "foo", "bar"},
"",
[]string{"foo", "bar"},
false,
},
{
"both env var and CLI",
[]string{testCommandName, "foo", "bar"},
"-foo bar",
[]string{"-foo", "bar", "foo", "bar"},
false,
},
{
"only env var",
[]string{testCommandName},
"-foo bar",
[]string{"-foo", "bar"},
false,
},
{
// this should fail gracefully, this is just testing
// that we don't panic with our slice arithmetic
"no command",
[]string{},
"-foo bar",
nil,
true,
},
{
"single quoted strings",
[]string{testCommandName, "foo"},
"-foo 'bar baz'",
[]string{"-foo", "bar baz", "foo"},
false,
},
{
"double quoted strings",
[]string{testCommandName, "foo"},
`-foo "bar baz"`,
[]string{"-foo", "bar baz", "foo"},
false,
},
{
"double quoted single quoted strings",
[]string{testCommandName, "foo"},
`-foo "'bar baz'"`,
[]string{"-foo", "'bar baz'", "foo"},
false,
},
}
for i, tc := range cases {
t.Run(fmt.Sprintf("%d-%s", i, tc.Name), func(t *testing.T) {
os.Unsetenv(EnvCLI)
// Set the env var value
if tc.Value != "" {
if err := os.Setenv(EnvCLI, 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 {
Args []string
}
func (c *testCommandCLI) Run(args []string) int {
c.Args = args
return 0
}
func (c *testCommandCLI) Synopsis() string { return "" }
func (c *testCommandCLI) Help() string { return "" }