make apply shutdown test completely deterministic

Add a shutdown hook to verify that a context has been correctly
cancelled, so we can remove the sleep and stop guessing.

Add a plan version of the shutdown test as well.

command/apply.go

@@ -183,6 +183,11 @@ func (c *ApplyCommand) Run(args []string) int {
 		// Cancel our context so we can start gracefully exiting
 		ctxCancel()
 
+		// notify tests that the command context was canceled
+		if testShutdownHook != nil {
+			testShutdownHook()
+		}
+
 		// Notify the user
 		c.Ui.Output(outputInterrupt)
 

command/apply_test.go

@@ -824,14 +824,20 @@ func TestApply_refresh(t *testing.T) {
 }
 
 func TestApply_shutdown(t *testing.T) {
-	stopped := false
+	cancelled := false
-	stopCh := make(chan struct{})
+	cancelDone := make(chan struct{})
-	stopReplyCh := make(chan struct{})
+	testShutdownHook = func() {
+		cancelled = true
+		close(cancelDone)
+	}
+	defer func() {
+		testShutdownHook = nil
+	}()
 
 	statePath := testTempFile(t)
-
 	p := testProvider()
 	shutdownCh := make(chan struct{})
+
 	ui := new(cli.MockUi)
 	c := &ApplyCommand{
 		Meta: Meta{
@@ -857,10 +863,10 @@ func TestApply_shutdown(t *testing.T) {
 		*terraform.InstanceInfo,
 		*terraform.InstanceState,
 		*terraform.InstanceDiff) (*terraform.InstanceState, error) {
-		if !stopped {
+
-			stopped = true
+		if !cancelled {
-			close(stopCh)
+			shutdownCh <- struct{}{}
-			<-stopReplyCh
+			<-cancelDone
 		}
 
 		return &terraform.InstanceState{
@@ -871,18 +877,6 @@ func TestApply_shutdown(t *testing.T) {
 		}, nil
 	}
 
-	go func() {
-		<-stopCh
-		shutdownCh <- struct{}{}
-
-		// This is really dirty, but we have no other way to assure that
-		// tf.Stop() has been called. This doesn't assure it either, but
-		// it makes it much more certain.
-		time.Sleep(50 * time.Millisecond)
-
-		close(stopReplyCh)
-	}()
-
 	args := []string{
 		"-state", statePath,
 		"-auto-approve",
@@ -896,6 +890,10 @@ func TestApply_shutdown(t *testing.T) {
 		t.Fatalf("err: %s", err)
 	}
 
+	if !cancelled {
+		t.Fatal("command not cancelled")
+	}
+
 	state := testStateRead(t, statePath)
 	if state == nil {
 		t.Fatal("state should not be nil")

command/meta.go

@@ -641,3 +641,7 @@ func isAutoVarFile(path string) bool {
 	return strings.HasSuffix(path, ".auto.tfvars") ||
 		strings.HasSuffix(path, ".auto.tfvars.json")
 }
+
+// testShutdownHook is used by tests to verify that a command context has been
+// canceled
+var testShutdownHook func()

command/plan.go

@@ -117,6 +117,12 @@ func (c *PlanCommand) Run(args []string) int {
 	case <-c.ShutdownCh:
 		// Cancel our context so we can start gracefully exiting
 		ctxCancel()
+
+		// notify tests that the command context was canceled
+		if testShutdownHook != nil {
+			testShutdownHook()
+		}
+
 		// Notify the user
 		c.Ui.Output(outputInterrupt)
 

command/plan_test.go

@@ -831,6 +831,56 @@ func TestPlan_detailedExitcode_emptyDiff(t *testing.T) {
 	}
 }
 
+func TestPlan_shutdown(t *testing.T) {
+	cancelled := false
+	cancelDone := make(chan struct{})
+	testShutdownHook = func() {
+		cancelled = true
+		close(cancelDone)
+	}
+	defer func() {
+		testShutdownHook = nil
+	}()
+
+	shutdownCh := make(chan struct{})
+	p := testProvider()
+	ui := new(cli.MockUi)
+	c := &PlanCommand{
+		Meta: Meta{
+			testingOverrides: metaOverridesForProvider(p),
+			Ui:               ui,
+			ShutdownCh:       shutdownCh,
+		},
+	}
+
+	p.DiffFn = func(
+		*terraform.InstanceInfo,
+		*terraform.InstanceState,
+		*terraform.ResourceConfig) (*terraform.InstanceDiff, error) {
+
+		if !cancelled {
+			shutdownCh <- struct{}{}
+			<-cancelDone
+		}
+
+		return &terraform.InstanceDiff{
+			Attributes: map[string]*terraform.ResourceAttrDiff{
+				"ami": &terraform.ResourceAttrDiff{
+					New: "bar",
+				},
+			},
+		}, nil
+	}
+
+	if code := c.Run([]string{testFixturePath("apply-shutdown")}); code != 0 {
+		t.Fatalf("bad: %d\n\n%s", code, ui.ErrorWriter.String())
+	}
+
+	if !cancelled {
+		t.Fatal("command not cancelled")
+	}
+}
+
 const planVarFile = `
 foo = "bar"
 `