diff --git a/commands.go b/commands.go
index 80cf878bf..eec2f520e 100644
--- a/commands.go
+++ b/commands.go
@@ -25,15 +25,7 @@ const (
 	OutputPrefix = "o:"
 )
 
-func init() {
-	Ui = &cli.PrefixedUi{
-		AskPrefix:    OutputPrefix,
-		OutputPrefix: OutputPrefix,
-		InfoPrefix:   OutputPrefix,
-		ErrorPrefix:  ErrorPrefix,
-		Ui:           &cli.BasicUi{Writer: os.Stdout},
-	}
-
+func initCommands(config *Config) {
 	var inAutomation bool
 	if v := os.Getenv(runningInAutomationEnvName); v != "" {
 		inAutomation = true
diff --git a/main.go b/main.go
index b088d255e..62e4da603 100644
--- a/main.go
+++ b/main.go
@@ -96,6 +96,16 @@ func realMain() int {
 	return wrappedMain()
 }
 
+func init() {
+	Ui = &cli.PrefixedUi{
+		AskPrefix:    OutputPrefix,
+		OutputPrefix: OutputPrefix,
+		InfoPrefix:   OutputPrefix,
+		ErrorPrefix:  ErrorPrefix,
+		Ui:           &cli.BasicUi{Writer: os.Stdout},
+	}
+}
+
 func wrappedMain() int {
 	// We always need to close the DebugInfo before we exit.
 	defer terraform.CloseDebugInfo()
@@ -128,6 +138,11 @@ func wrappedMain() int {
 		config = *config.Merge(usrcfg)
 	}
 
+	// In tests, Commands may already be set to provide mock commands
+	if Commands == nil {
+		initCommands(&config)
+	}
+
 	// Run checkpoint
 	go runCheckpoint(&config)
 
diff --git a/main_test.go b/main_test.go
index e5500e8ef..406195e8b 100644
--- a/main_test.go
+++ b/main_test.go
@@ -18,9 +18,12 @@ func TestMain_cliArgsFromEnv(t *testing.T) {
 	defer func() { os.Args = oldArgs }()
 
 	// Setup test command and restore that
+	Commands = make(map[string]cli.CommandFactory)
+	defer func() {
+		Commands = nil
+	}()
 	testCommandName := "unit-test-cli-args"
 	testCommand := &testCommandCLI{}
-	defer func() { delete(Commands, testCommandName) }()
 	Commands[testCommandName] = func() (cli.Command, error) {
 		return testCommand, nil
 	}
@@ -150,6 +153,12 @@ func TestMain_cliArgsFromEnvAdvanced(t *testing.T) {
 	oldArgs := os.Args
 	defer func() { os.Args = oldArgs }()
 
+	// Setup test command and restore that
+	Commands = make(map[string]cli.CommandFactory)
+	defer func() {
+		Commands = nil
+	}()
+
 	cases := []struct {
 		Name     string
 		Command  string
@@ -230,7 +239,7 @@ func TestMain_cliArgsFromEnvAdvanced(t *testing.T) {
 			testCommand.Args = nil
 			exit := wrappedMain()
 			if (exit != 0) != tc.Err {
-				t.Fatalf("bad: %d", exit)
+				t.Fatalf("unexpected exit status %d; want 0", exit)
 			}
 			if tc.Err {
 				return