diff --git a/command/apply.go b/command/apply.go
index 3c3d2a189..84b16ce17 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 2707f293c..3ae4d7adb 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 729020b2b..27f7765f9 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 f6f4f8f2f..bef831e16 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 152486c7f..04390ab5b 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{