helper/schema: properly put "id" into attributes

helper/schema/resource.go

@@ -81,17 +81,7 @@ func (r *Resource) Apply(
 		err = r.Update(data, meta)
 	}
 
-	// Always set the ID attribute if it is set. We also always collapse
+	return data.State(), err
-	// the state since even partial states need to be returned.
-	state := data.State()
-	if state.ID != "" {
-		if state.Attributes == nil {
-			state.Attributes = make(map[string]string)
-		}
-		state.Attributes["id"] = state.ID
-	}
-
-	return state, err
 }
 
 // Diff returns a diff of this resource and is API compatible with the

helper/schema/resource_data.go

@@ -129,6 +129,10 @@ func (d *ResourceData) State() *terraform.ResourceState {
 	result.Attributes = d.stateObject("", d.schema)
 	result.Dependencies = d.Dependencies()
 
+	if v := d.Id(); v != "" {
+		result.Attributes["id"] = d.Id()
+	}
+
 	return &result
 }
 

helper/schema/resource_data_test.go

@@ -1309,6 +1309,43 @@ func TestResourceDataState(t *testing.T) {
 				Attributes: map[string]string{},
 			},
 		},
+
+		// Basic state with other keys
+		{
+			Schema: map[string]*Schema{
+				"availability_zone": &Schema{
+					Type:     TypeString,
+					Optional: true,
+					Computed: true,
+					ForceNew: true,
+				},
+			},
+
+			State: &terraform.ResourceState{
+				ID: "bar",
+				Attributes: map[string]string{
+					"id": "bar",
+				},
+			},
+
+			Diff: &terraform.ResourceDiff{
+				Attributes: map[string]*terraform.ResourceAttrDiff{
+					"availability_zone": &terraform.ResourceAttrDiff{
+						Old:         "",
+						New:         "foo",
+						RequiresNew: true,
+					},
+				},
+			},
+
+			Result: &terraform.ResourceState{
+				ID: "bar",
+				Attributes: map[string]string{
+					"id":                "bar",
+					"availability_zone": "foo",
+				},
+			},
+		},
 	}
 
 	for i, tc := range cases {

helper/schema/resource_test.go

@@ -129,6 +129,7 @@ func TestResourceApply_destroyPartial(t *testing.T) {
 	expected := &terraform.ResourceState{
 		ID: "bar",
 		Attributes: map[string]string{
+			"id":  "bar",
 			"foo": "42",
 		},
 	}
@@ -291,6 +292,7 @@ func TestResourceRefresh(t *testing.T) {
 	expected := &terraform.ResourceState{
 		ID: "bar",
 		Attributes: map[string]string{
+			"id":  "bar",
 			"foo": "13",
 		},
 	}