From 059dcf1e252427ccfa1df5c3d94bfbeed0e8938d Mon Sep 17 00:00:00 2001
From: James Bardin <j.bardin@gmail.com>
Date: Tue, 14 May 2019 18:00:01 -0400
Subject: [PATCH] implement UpgradeResourceState in the mock provider

The default is a Noop that simply encodes the provided state
---
 terraform/provider_mock.go | 28 +++++++++++++++++++++++++++-
 1 file changed, 27 insertions(+), 1 deletion(-)

diff --git a/terraform/provider_mock.go b/terraform/provider_mock.go
index 8c5824089..4ae346d7d 100644
--- a/terraform/provider_mock.go
+++ b/terraform/provider_mock.go
@@ -6,6 +6,7 @@ import (
 	"sync"
 
 	"github.com/zclconf/go-cty/cty"
+	ctyjson "github.com/zclconf/go-cty/cty/json"
 
 	"github.com/hashicorp/terraform/config"
 	"github.com/hashicorp/terraform/config/hcl2shim"
@@ -191,6 +192,10 @@ func (p *MockProvider) UpgradeResourceState(r providers.UpgradeResourceStateRequ
 	p.Lock()
 	defer p.Unlock()
 
+	schemas := p.getSchema()
+	schema := schemas.ResourceTypes[r.TypeName]
+	schemaType := schema.Block.ImpliedType()
+
 	p.UpgradeResourceStateCalled = true
 	p.UpgradeResourceStateRequest = r
 
@@ -198,7 +203,28 @@ func (p *MockProvider) UpgradeResourceState(r providers.UpgradeResourceStateRequ
 		return p.UpgradeResourceStateFn(r)
 	}
 
-	return p.UpgradeResourceStateResponse
+	resp := p.UpgradeResourceStateResponse
+
+	if resp.UpgradedState == cty.NilVal {
+		switch {
+		case r.RawStateFlatmap != nil:
+			v, err := hcl2shim.HCL2ValueFromFlatmap(r.RawStateFlatmap, schemaType)
+			if err != nil {
+				resp.Diagnostics = resp.Diagnostics.Append(err)
+				return resp
+			}
+			resp.UpgradedState = v
+		case len(r.RawStateJSON) > 0:
+			v, err := ctyjson.Unmarshal(r.RawStateJSON, schemaType)
+
+			if err != nil {
+				resp.Diagnostics = resp.Diagnostics.Append(err)
+				return resp
+			}
+			resp.UpgradedState = v
+		}
+	}
+	return resp
 }
 
 func (p *MockProvider) Configure(r providers.ConfigureRequest) providers.ConfigureResponse {