From 484d67028a48b404931e18900361caed061978bc Mon Sep 17 00:00:00 2001
From: James Bardin <j.bardin@gmail.com>
Date: Tue, 4 Dec 2018 16:00:43 -0500
Subject: [PATCH] handle shim errors in provider tests

These should be rare, and even though it's likely a shim bug, the error
is probably easier for provider developers to deal with than the
panic
---
 helper/resource/testing_config.go | 24 ++++++++++++++++++++----
 1 file changed, 20 insertions(+), 4 deletions(-)

diff --git a/helper/resource/testing_config.go b/helper/resource/testing_config.go
index 9bf04cb59..6e995e158 100644
--- a/helper/resource/testing_config.go
+++ b/helper/resource/testing_config.go
@@ -45,7 +45,11 @@ func testStep(opts terraform.ContextOpts, state *terraform.State, step TestStep,
 
 	// Build the context
 	opts.Config = cfg
-	opts.State = terraform.MustShimLegacyState(state)
+	opts.State, err = terraform.ShimLegacyState(state)
+	if err != nil {
+		return nil, err
+	}
+
 	opts.Destroy = step.Destroy
 	ctx, stepDiags := terraform.NewContext(&opts)
 	if stepDiags.HasErrors() {
@@ -61,7 +65,11 @@ func testStep(opts terraform.ContextOpts, state *terraform.State, step TestStep,
 
 	// Refresh!
 	newState, stepDiags := ctx.Refresh()
-	state = mustShimNewState(newState, schemas)
+	// shim the state first so the test can check the state on errors
+	state, err = shimNewState(newState, schemas)
+	if err != nil {
+		return nil, err
+	}
 	if stepDiags.HasErrors() {
 		return state, fmt.Errorf("Error refreshing: %s", stepDiags.Err())
 	}
@@ -83,7 +91,11 @@ func testStep(opts terraform.ContextOpts, state *terraform.State, step TestStep,
 
 		// Apply the diff, creating real resources.
 		newState, stepDiags = ctx.Apply()
-		state = mustShimNewState(newState, schemas)
+		// shim the state first so the test can check the state on errors
+		state, err = shimNewState(newState, schemas)
+		if err != nil {
+			return nil, err
+		}
 		if stepDiags.HasErrors() {
 			return state, fmt.Errorf("Error applying: %s", stepDiags.Err())
 		}
@@ -123,7 +135,11 @@ func testStep(opts terraform.ContextOpts, state *terraform.State, step TestStep,
 		if stepDiags.HasErrors() {
 			return state, fmt.Errorf("Error on follow-up refresh: %s", stepDiags.Err())
 		}
-		state = mustShimNewState(newState, schemas)
+
+		state, err = shimNewState(newState, schemas)
+		if err != nil {
+			return nil, err
+		}
 	}
 	if p, stepDiags = ctx.Plan(); stepDiags.HasErrors() {
 		return state, fmt.Errorf("Error on second follow-up plan: %s", stepDiags.Err())