diff --git a/command/apply_test.go b/command/apply_test.go
index da755d329..c27d184d8 100644
--- a/command/apply_test.go
+++ b/command/apply_test.go
@@ -341,10 +341,7 @@ func TestApply_error(t *testing.T) {
 
 	var lock sync.Mutex
 	errored := false
-	p.ApplyFn = func(
-		info *terraform.InstanceInfo,
-		s *terraform.InstanceState,
-		d *terraform.InstanceDiff) (*terraform.InstanceState, error) {
+	p.ApplyFn = func(info *terraform.InstanceInfo, s *terraform.InstanceState, d *terraform.InstanceDiff) (*terraform.InstanceState, error) {
 		lock.Lock()
 		defer lock.Unlock()
 
@@ -353,19 +350,34 @@ func TestApply_error(t *testing.T) {
 			return nil, fmt.Errorf("error")
 		}
 
-		return &terraform.InstanceState{ID: "foo"}, nil
+		newState := &terraform.InstanceState{
+			ID:         "foo",
+			Attributes: map[string]string{},
+		}
+		newState.Attributes["id"] = newState.ID
+		if ad, ok := d.Attributes["ami"]; ok {
+			newState.Attributes["ami"] = ad.New
+		}
+		if ad, ok := d.Attributes["error"]; ok {
+			newState.Attributes["error"] = ad.New
+		}
+		return newState, nil
 	}
-	p.DiffFn = func(
-		*terraform.InstanceInfo,
-		*terraform.InstanceState,
-		*terraform.ResourceConfig) (*terraform.InstanceDiff, error) {
-		return &terraform.InstanceDiff{
-			Attributes: map[string]*terraform.ResourceAttrDiff{
-				"ami": &terraform.ResourceAttrDiff{
-					New: "bar",
-				},
-			},
-		}, nil
+	p.DiffFn = func(info *terraform.InstanceInfo, s *terraform.InstanceState, rc *terraform.ResourceConfig) (*terraform.InstanceDiff, error) {
+		ret := &terraform.InstanceDiff{
+			Attributes: map[string]*terraform.ResourceAttrDiff{},
+		}
+		if new, ok := rc.Get("ami"); ok {
+			ret.Attributes["ami"] = &terraform.ResourceAttrDiff{
+				New: new.(string),
+			}
+		}
+		if new, ok := rc.Get("error"); ok {
+			ret.Attributes["error"] = &terraform.ResourceAttrDiff{
+				New: fmt.Sprintf("%t", new.(bool)),
+			}
+		}
+		return ret, nil
 	}
 	p.GetSchemaReturn = &terraform.ProviderSchema{
 		ResourceTypes: map[string]*configschema.Block{
@@ -1102,6 +1114,7 @@ func TestApply_vars(t *testing.T) {
 		},
 	}
 
+	actual := ""
 	p.GetSchemaReturn = &terraform.ProviderSchema{
 		ResourceTypes: map[string]*configschema.Block{
 			"test_instance": {
@@ -1116,17 +1129,11 @@ func TestApply_vars(t *testing.T) {
 			NewState: req.PlannedState,
 		}
 	}
-
-	actual := ""
-	p.DiffFn = func(
-		info *terraform.InstanceInfo,
-		s *terraform.InstanceState,
-		c *terraform.ResourceConfig) (*terraform.InstanceDiff, error) {
-		if v, ok := c.Config["value"]; ok {
-			actual = v.(string)
+	p.PlanResourceChangeFn = func(req providers.PlanResourceChangeRequest) providers.PlanResourceChangeResponse {
+		actual = req.ProposedNewState.GetAttr("value").AsString()
+		return providers.PlanResourceChangeResponse{
+			PlannedState: req.ProposedNewState,
 		}
-
-		return &terraform.InstanceDiff{}, nil
 	}
 
 	args := []string{
@@ -1161,6 +1168,7 @@ func TestApply_varFile(t *testing.T) {
 		},
 	}
 
+	actual := ""
 	p.GetSchemaReturn = &terraform.ProviderSchema{
 		ResourceTypes: map[string]*configschema.Block{
 			"test_instance": {
@@ -1175,17 +1183,11 @@ func TestApply_varFile(t *testing.T) {
 			NewState: req.PlannedState,
 		}
 	}
-
-	actual := ""
-	p.DiffFn = func(
-		info *terraform.InstanceInfo,
-		s *terraform.InstanceState,
-		c *terraform.ResourceConfig) (*terraform.InstanceDiff, error) {
-		if v, ok := c.Config["value"]; ok {
-			actual = v.(string)
+	p.PlanResourceChangeFn = func(req providers.PlanResourceChangeRequest) providers.PlanResourceChangeResponse {
+		actual = req.ProposedNewState.GetAttr("value").AsString()
+		return providers.PlanResourceChangeResponse{
+			PlannedState: req.ProposedNewState,
 		}
-
-		return &terraform.InstanceDiff{}, nil
 	}
 
 	args := []string{
@@ -1230,6 +1232,7 @@ func TestApply_varFileDefault(t *testing.T) {
 		},
 	}
 
+	actual := ""
 	p.GetSchemaReturn = &terraform.ProviderSchema{
 		ResourceTypes: map[string]*configschema.Block{
 			"test_instance": {
@@ -1244,17 +1247,11 @@ func TestApply_varFileDefault(t *testing.T) {
 			NewState: req.PlannedState,
 		}
 	}
-
-	actual := ""
-	p.DiffFn = func(
-		info *terraform.InstanceInfo,
-		s *terraform.InstanceState,
-		c *terraform.ResourceConfig) (*terraform.InstanceDiff, error) {
-		if v, ok := c.Config["value"]; ok {
-			actual = v.(string)
+	p.PlanResourceChangeFn = func(req providers.PlanResourceChangeRequest) providers.PlanResourceChangeResponse {
+		actual = req.ProposedNewState.GetAttr("value").AsString()
+		return providers.PlanResourceChangeResponse{
+			PlannedState: req.ProposedNewState,
 		}
-
-		return &terraform.InstanceDiff{}, nil
 	}
 
 	args := []string{
@@ -1298,6 +1295,7 @@ func TestApply_varFileDefaultJSON(t *testing.T) {
 		},
 	}
 
+	actual := ""
 	p.GetSchemaReturn = &terraform.ProviderSchema{
 		ResourceTypes: map[string]*configschema.Block{
 			"test_instance": {
@@ -1312,17 +1310,11 @@ func TestApply_varFileDefaultJSON(t *testing.T) {
 			NewState: req.PlannedState,
 		}
 	}
-
-	actual := ""
-	p.DiffFn = func(
-		info *terraform.InstanceInfo,
-		s *terraform.InstanceState,
-		c *terraform.ResourceConfig) (*terraform.InstanceDiff, error) {
-		if v, ok := c.Config["value"]; ok {
-			actual = v.(string)
+	p.PlanResourceChangeFn = func(req providers.PlanResourceChangeRequest) providers.PlanResourceChangeResponse {
+		actual = req.ProposedNewState.GetAttr("value").AsString()
+		return providers.PlanResourceChangeResponse{
+			PlannedState: req.ProposedNewState,
 		}
-
-		return &terraform.InstanceDiff{}, nil
 	}
 
 	args := []string{