From 2941ed464ca09f804c05b22149a6f74c152bcd54 Mon Sep 17 00:00:00 2001
From: James Bardin <j.bardin@gmail.com>
Date: Sat, 2 Dec 2017 22:36:43 -0500
Subject: [PATCH] replace the testShutdownHook with a check for Stop

Now that the local backend can be cancelled during plan and refresh, we
don't really need the testShutdownHook. Simplify the tests by just
checking for Stop being called on the provider.
---
 command/apply.go      |  6 +-----
 command/apply_test.go | 23 +++++++++--------------
 command/meta.go       |  4 ----
 command/plan.go       |  5 -----
 command/plan_test.go  | 17 ++++++++---------
 5 files changed, 18 insertions(+), 37 deletions(-)

diff --git a/command/apply.go b/command/apply.go
index 3c3d2a1894..84b16ce17f 100644
--- a/command/apply.go
+++ b/command/apply.go
@@ -171,6 +171,7 @@ func (c *ApplyCommand) Run(args []string) int {
 	// Perform the operation
 	ctx, ctxCancel := context.WithCancel(context.Background())
 	defer ctxCancel()
+
 	op, err := b.Operation(ctx, opReq)
 	if err != nil {
 		c.Ui.Error(fmt.Sprintf("Error starting operation: %s", err))
@@ -183,11 +184,6 @@ 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)
 
diff --git a/command/apply_test.go b/command/apply_test.go
index 2707f293c3..3ae4d7adb1 100644
--- a/command/apply_test.go
+++ b/command/apply_test.go
@@ -825,14 +825,7 @@ func TestApply_refresh(t *testing.T) {
 
 func TestApply_shutdown(t *testing.T) {
 	cancelled := false
-	cancelDone := make(chan struct{})
-	testShutdownHook = func() {
-		cancelled = true
-		close(cancelDone)
-	}
-	defer func() {
-		testShutdownHook = nil
-	}()
+	stopped := make(chan struct{})
 
 	statePath := testTempFile(t)
 	p := testProvider()
@@ -847,6 +840,12 @@ func TestApply_shutdown(t *testing.T) {
 		},
 	}
 
+	p.StopFn = func() error {
+		close(stopped)
+		cancelled = true
+		return nil
+	}
+
 	p.DiffFn = func(
 		*terraform.InstanceInfo,
 		*terraform.InstanceState,
@@ -864,11 +863,11 @@ func TestApply_shutdown(t *testing.T) {
 		*terraform.InstanceState,
 		*terraform.InstanceDiff) (*terraform.InstanceState, error) {
 
+		// only cancel once
 		if !cancelled {
 			shutdownCh <- struct{}{}
-			<-cancelDone
+			<-stopped
 		}
-
 		return &terraform.InstanceState{
 			ID: "foo",
 			Attributes: map[string]string{
@@ -898,10 +897,6 @@ func TestApply_shutdown(t *testing.T) {
 	if state == nil {
 		t.Fatal("state should not be nil")
 	}
-
-	if len(state.RootModule().Resources) != 1 {
-		t.Fatalf("bad: %d", len(state.RootModule().Resources))
-	}
 }
 
 func TestApply_state(t *testing.T) {
diff --git a/command/meta.go b/command/meta.go
index 729020b2ba..27f7765f95 100644
--- a/command/meta.go
+++ b/command/meta.go
@@ -641,7 +641,3 @@ 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()
diff --git a/command/plan.go b/command/plan.go
index f6f4f8f2f7..bef831e16d 100644
--- a/command/plan.go
+++ b/command/plan.go
@@ -118,11 +118,6 @@ func (c *PlanCommand) 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)
 
diff --git a/command/plan_test.go b/command/plan_test.go
index 152486c7ff..04390ab5bc 100644
--- a/command/plan_test.go
+++ b/command/plan_test.go
@@ -833,14 +833,7 @@ 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
-	}()
+	stopped := make(chan struct{})
 
 	shutdownCh := make(chan struct{})
 	p := testProvider()
@@ -853,6 +846,12 @@ func TestPlan_shutdown(t *testing.T) {
 		},
 	}
 
+	p.StopFn = func() error {
+		close(stopped)
+		cancelled = true
+		return nil
+	}
+
 	p.DiffFn = func(
 		*terraform.InstanceInfo,
 		*terraform.InstanceState,
@@ -860,7 +859,7 @@ func TestPlan_shutdown(t *testing.T) {
 
 		if !cancelled {
 			shutdownCh <- struct{}{}
-			<-cancelDone
+			<-stopped
 		}
 
 		return &terraform.InstanceDiff{