Skip to content

Commit df93e51

Browse files
mitchellhmitchellh
committed
allow targeted TF_CLI_ARGS_x
1 parent 53796fc commit df93e51

2 files changed

Lines changed: 142 additions & 33 deletions

File tree

‎main.go‎

Lines changed: 60 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import (
77
"log"
88
"os"
99
"runtime"
10+
"strings"
1011
"sync"
1112

1213
"github.com/hashicorp/go-plugin"
@@ -138,39 +139,27 @@ func wrappedMain() int {
138139
// Get the command line args.
139140
args := os.Args[1:]
140141

141-
// Prefix the args with any args from the EnvCLI
142-
if v := os.Getenv(EnvCLI); v != "" {
143-
log.Printf("[INFO] %s value: %q", EnvCLI, v)
144-
extra, err := shellwords.Parse(v)
145-
if err != nil {
146-
Ui.Error(fmt.Sprintf(
147-
"Error parsing extra CLI args from %s: %s",
148-
EnvCLI, err))
149-
return 1
150-
}
151-
152-
// Find the index to place the flags. We put them exactly
153-
// after the first non-flag arg.
154-
idx := -1
155-
for i, v := range args {
156-
if len(v) > 0 && v[0] != '-' {
157-
idx = i
158-
break
159-
}
160-
}
161-
162-
// idx points to the exact arg that isn't a flag. We increment
163-
// by one so that all the copying below expects idx to be the
164-
// insertion point.
165-
idx++
142+
// Build the CLI so far, we do this so we can query the subcommand.
143+
cliRunner := &cli.CLI{
144+
Args: args,
145+
Commands: Commands,
146+
HelpFunc: helpFunc,
147+
HelpWriter: os.Stdout,
148+
}
166149

167-
// Copy the args
168-
newArgs := make([]string, len(args)+len(extra))
169-
copy(newArgs, args[:idx])
170-
copy(newArgs[idx:], extra)
171-
copy(newArgs[len(extra)+idx:], args[idx:])
172-
args = newArgs
150+
// Prefix the args with any args from the EnvCLI
151+
args, err = mergeEnvArgs(EnvCLI, args)
152+
if err != nil {
153+
Ui.Error(err.Error())
154+
return 1
155+
}
173156

157+
// Prefix the args with any args from the EnvCLI targeting this command
158+
suffix := strings.Replace(cliRunner.Subcommand(), "-", "_", -1)
159+
args, err = mergeEnvArgs(fmt.Sprintf("%s_%s", EnvCLI, suffix), args)
160+
if err != nil {
161+
Ui.Error(err.Error())
162+
return 1
174163
}
175164

176165
// We shortcut "--version" and "-v" to just show the version
@@ -184,8 +173,9 @@ func wrappedMain() int {
184173
}
185174
}
186175

176+
// Rebuild the CLI with any modified args.
187177
log.Printf("[INFO] CLI command args: %#v", args)
188-
cli := &cli.CLI{
178+
cliRunner = &cli.CLI{
189179
Args: args,
190180
Commands: Commands,
191181
HelpFunc: helpFunc,
@@ -196,7 +186,7 @@ func wrappedMain() int {
196186
ContextOpts.Providers = config.ProviderFactories()
197187
ContextOpts.Provisioners = config.ProvisionerFactories()
198188

199-
exitCode, err := cli.Run()
189+
exitCode, err := cliRunner.Run()
200190
if err != nil {
201191
Ui.Error(fmt.Sprintf("Error executing CLI: %s", err.Error()))
202192
return 1
@@ -284,3 +274,40 @@ func copyOutput(r io.Reader, doneCh chan<- struct{}) {
284274

285275
wg.Wait()
286276
}
277+
278+
func mergeEnvArgs(envName string, args []string) ([]string, error) {
279+
v := os.Getenv(envName)
280+
if v == "" {
281+
return args, nil
282+
}
283+
284+
log.Printf("[INFO] %s value: %q", envName, v)
285+
extra, err := shellwords.Parse(v)
286+
if err != nil {
287+
return nil, fmt.Errorf(
288+
"Error parsing extra CLI args from %s: %s",
289+
envName, err)
290+
}
291+
292+
// Find the index to place the flags. We put them exactly
293+
// after the first non-flag arg.
294+
idx := -1
295+
for i, v := range args {
296+
if len(v) > 0 && v[0] != '-' {
297+
idx = i
298+
break
299+
}
300+
}
301+
302+
// idx points to the exact arg that isn't a flag. We increment
303+
// by one so that all the copying below expects idx to be the
304+
// insertion point.
305+
idx++
306+
307+
// Copy the args
308+
newArgs := make([]string, len(args)+len(extra))
309+
copy(newArgs, args[:idx])
310+
copy(newArgs[idx:], extra)
311+
copy(newArgs[len(extra)+idx:], args[idx:])
312+
return newArgs, nil
313+
}

‎main_test.go‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -110,6 +110,7 @@ func TestMain_cliArgsFromEnv(t *testing.T) {
110110
for i, tc := range cases {
111111
t.Run(fmt.Sprintf("%d-%s", i, tc.Name), func(t *testing.T) {
112112
os.Unsetenv(EnvCLI)
113+
defer os.Unsetenv(EnvCLI)
113114

114115
// Set the env var value
115116
if tc.Value != "" {
@@ -142,6 +143,87 @@ func TestMain_cliArgsFromEnv(t *testing.T) {
142143
}
143144
}
144145

146+
// This test just has more options than the test above. Use this for
147+
// more control over behavior at the expense of more complex test structures.
148+
func TestMain_cliArgsFromEnvAdvanced(t *testing.T) {
149+
// Restore original CLI args
150+
oldArgs := os.Args
151+
defer func() { os.Args = oldArgs }()
152+
153+
cases := []struct {
154+
Name string
155+
Command string
156+
EnvVar string
157+
Args []string
158+
Value string
159+
Expected []string
160+
Err bool
161+
}{
162+
{
163+
"targeted to another command",
164+
"command",
165+
EnvCLI + "_foo",
166+
[]string{"command", "foo", "bar"},
167+
"-flag",
168+
[]string{"foo", "bar"},
169+
false,
170+
},
171+
172+
{
173+
"targeted to this command",
174+
"command",
175+
EnvCLI + "_command",
176+
[]string{"command", "foo", "bar"},
177+
"-flag",
178+
[]string{"-flag", "foo", "bar"},
179+
false,
180+
},
181+
}
182+
183+
for i, tc := range cases {
184+
t.Run(fmt.Sprintf("%d-%s", i, tc.Name), func(t *testing.T) {
185+
// Setup test command and restore that
186+
testCommandName := tc.Command
187+
testCommand := &testCommandCLI{}
188+
defer func() { delete(Commands, testCommandName) }()
189+
Commands[testCommandName] = func() (cli.Command, error) {
190+
return testCommand, nil
191+
}
192+
193+
os.Unsetenv(tc.EnvVar)
194+
defer os.Unsetenv(tc.EnvVar)
195+
196+
// Set the env var value
197+
if tc.Value != "" {
198+
if err := os.Setenv(tc.EnvVar, tc.Value); err != nil {
199+
t.Fatalf("err: %s", err)
200+
}
201+
}
202+
203+
// Setup the args
204+
args := make([]string, len(tc.Args)+1)
205+
args[0] = oldArgs[0] // process name
206+
copy(args[1:], tc.Args)
207+
208+
// Run it!
209+
os.Args = args
210+
testCommand.Args = nil
211+
exit := wrappedMain()
212+
if (exit != 0) != tc.Err {
213+
t.Fatalf("bad: %d", exit)
214+
}
215+
if tc.Err {
216+
return
217+
}
218+
219+
// Verify
220+
if !reflect.DeepEqual(testCommand.Args, tc.Expected) {
221+
t.Fatalf("bad: %#v", testCommand.Args)
222+
}
223+
})
224+
}
225+
}
226+
145227
type testCommandCLI struct {
146228
Args []string
147229
}

0 commit comments

Comments
 (0)