From 50026a6d5cc62a73ddd26da6a1619db6af7a1e64 Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Fri, 22 Aug 2014 08:57:44 -0700
Subject: [PATCH] helper/schema: When having a StateFunc, make sure NewExtra
 contains original

---
 helper/schema/resource_data.go      | 11 +++++++----
 helper/schema/resource_data_test.go | 26 ++++++++++++++++++++++++++
 helper/schema/schema.go             |  6 ++++++
 helper/schema/schema_test.go        |  5 +++--
 4 files changed, 42 insertions(+), 6 deletions(-)

diff --git a/helper/schema/resource_data.go b/helper/schema/resource_data.go
index e716e288f..f79d1b4a3 100644
--- a/helper/schema/resource_data.go
+++ b/helper/schema/resource_data.go
@@ -201,10 +201,6 @@ func (d *ResourceData) diffChange(k string) (interface{}, interface{}, bool) {
 		n.Value = nil
 	}
 
-	if n.Exists && n.Schema.StateFunc != nil {
-		n.Value = n.Schema.StateFunc(n.Value)
-	}
-
 	// Return the old, new, and whether there is a change
 	return o.Value, n.Value, !reflect.DeepEqual(o.Value, n.Value)
 }
@@ -512,6 +508,13 @@ func (d *ResourceData) getPrimitive(
 		attrD, ok := d.diff.Attributes[k]
 		if ok && !attrD.NewComputed {
 			result = attrD.New
+			if attrD.NewExtra != nil {
+				err := mapstructure.WeakDecode(attrD.NewExtra, &result)
+				if err != nil {
+					panic(err)
+				}
+			}
+
 			resultSet = true
 		}
 	}
diff --git a/helper/schema/resource_data_test.go b/helper/schema/resource_data_test.go
index ba07eca50..9fbb4c9d2 100644
--- a/helper/schema/resource_data_test.go
+++ b/helper/schema/resource_data_test.go
@@ -68,6 +68,32 @@ func TestResourceDataGet(t *testing.T) {
 			Value: "foo",
 		},
 
+		{
+			Schema: map[string]*Schema{
+				"availability_zone": &Schema{
+					Type:     TypeString,
+					Optional: true,
+					Computed: true,
+					ForceNew: true,
+				},
+			},
+
+			State: nil,
+
+			Diff: &terraform.ResourceDiff{
+				Attributes: map[string]*terraform.ResourceAttrDiff{
+					"availability_zone": &terraform.ResourceAttrDiff{
+						Old:         "",
+						New:         "foo!",
+						NewExtra:    "foo",
+					},
+				},
+			},
+
+			Key:   "availability_zone",
+			Value: "foo",
+		},
+
 		{
 			Schema: map[string]*Schema{
 				"availability_zone": &Schema{
diff --git a/helper/schema/schema.go b/helper/schema/schema.go
index 23f5c34fa..515b080d2 100644
--- a/helper/schema/schema.go
+++ b/helper/schema/schema.go
@@ -386,8 +386,13 @@ func (m schemaMap) diffString(
 	schema *Schema,
 	diff *terraform.ResourceDiff,
 	d *ResourceData) error {
+	var originalN interface{}
 	var os, ns string
 	o, n, _ := d.diffChange(k)
+	if schema.StateFunc != nil {
+		originalN = n
+		n = schema.StateFunc(n)
+	}
 	if err := mapstructure.WeakDecode(o, &os); err != nil {
 		return fmt.Errorf("%s: %s", k, err)
 	}
@@ -411,6 +416,7 @@ func (m schemaMap) diffString(
 	diff.Attributes[k] = schema.finalizeDiff(&terraform.ResourceAttrDiff{
 		Old:        os,
 		New:        ns,
+		NewExtra:   originalN,
 		NewRemoved: removed,
 	})
 
diff --git a/helper/schema/schema_test.go b/helper/schema/schema_test.go
index b457d4f89..ac1f827e5 100644
--- a/helper/schema/schema_test.go
+++ b/helper/schema/schema_test.go
@@ -98,8 +98,9 @@ func TestSchemaMap_Diff(t *testing.T) {
 			Diff: &terraform.ResourceDiff{
 				Attributes: map[string]*terraform.ResourceAttrDiff{
 					"availability_zone": &terraform.ResourceAttrDiff{
-						Old: "",
-						New: "foo!",
+						Old:      "",
+						New:      "foo!",
+						NewExtra: "foo",
 					},
 				},
 			},