From 032cd5f4426b9d7fb2e5411416da7b18dbd22351 Mon Sep 17 00:00:00 2001
From: Daniel Banck <dbanck@users.noreply.github.com>
Date: Tue, 29 Apr 2025 17:44:46 +0200
Subject: [PATCH] Fix version in resource identity upgrade request (#36940)

---
 internal/terraform/context_refresh_test.go   | 95 ++++++++++++++++++++
 internal/terraform/upgrade_resource_state.go |  2 +-
 2 files changed, 96 insertions(+), 1 deletion(-)

diff --git a/internal/terraform/context_refresh_test.go b/internal/terraform/context_refresh_test.go
index 3eade8d4f1..f9f9c2cea3 100644
--- a/internal/terraform/context_refresh_test.go
+++ b/internal/terraform/context_refresh_test.go
@@ -1691,3 +1691,98 @@ resource "test_resource" "foo" {
 		t.Fatalf("invalid state\nexpected: %s\ngot: %s\n", expected, jsonState)
 	}
 }
+
+func TestContext2Refresh_identityUpgradeJSON(t *testing.T) {
+	m := testModule(t, "refresh-schema-upgrade")
+	p := testProvider("test")
+	p.GetProviderSchemaResponse = getProviderSchemaResponseFromProviderSchema(&providerSchema{
+		ResourceTypes: map[string]*configschema.Block{
+			"test_thing": {
+				Attributes: map[string]*configschema.Attribute{
+					"id": {
+						Type:     cty.String,
+						Optional: true,
+					},
+				},
+			},
+		},
+		IdentityTypes: map[string]*configschema.Object{
+			"test_thing": {
+				Attributes: map[string]*configschema.Attribute{
+					"name": {
+						Type:     cty.String,
+						Required: true,
+					},
+				},
+				Nesting: configschema.NestingSingle,
+			},
+		},
+		IdentityTypeSchemaVersions: map[string]uint64{
+			"test_thing": 5,
+		},
+	})
+	p.UpgradeResourceIdentityResponse = &providers.UpgradeResourceIdentityResponse{
+		UpgradedIdentity: cty.ObjectVal(map[string]cty.Value{
+			"name": cty.StringVal("foo"),
+		}),
+	}
+
+	s := states.BuildState(func(s *states.SyncState) {
+		s.SetResourceInstanceCurrent(
+			addrs.Resource{
+				Mode: addrs.ManagedResourceMode,
+				Type: "test_thing",
+				Name: "bar",
+			}.Instance(addrs.NoKey).Absolute(addrs.RootModuleInstance),
+			&states.ResourceInstanceObjectSrc{
+				Status:                states.ObjectReady,
+				SchemaVersion:         0,
+				AttrsJSON:             []byte(`{"id":"foo"}`),
+				IdentitySchemaVersion: 3,
+				IdentityJSON:          []byte(`{"id":"foo"}`),
+			},
+			addrs.AbsProviderConfig{
+				Provider: addrs.NewDefaultProvider("test"),
+				Module:   addrs.RootModule,
+			},
+		)
+	})
+
+	ctx := testContext2(t, &ContextOpts{
+		Providers: map[addrs.Provider]providers.Factory{
+			addrs.NewDefaultProvider("test"): testProviderFuncFixed(p),
+		},
+	})
+
+	state, diags := ctx.Refresh(m, s, &PlanOpts{Mode: plans.NormalMode})
+	if diags.HasErrors() {
+		t.Fatal(diags.Err())
+	}
+
+	{
+		got := p.UpgradeResourceIdentityRequest
+		want := providers.UpgradeResourceIdentityRequest{
+			TypeName:        "test_thing",
+			Version:         3,
+			RawIdentityJSON: []byte(`{"id":"foo"}`),
+		}
+		if !cmp.Equal(got, want) {
+			t.Errorf("wrong identity upgrade request\n%s", cmp.Diff(want, got))
+		}
+	}
+
+	addr := mustResourceInstanceAddr("test_thing.bar")
+	res := state.ResourceInstance(addr)
+	if res == nil {
+		t.Fatalf("no resource in state for %s", addr)
+	}
+
+	expectedIdentity := `{"name":"foo"}`
+	if string(res.Current.IdentityJSON) != expectedIdentity {
+		t.Fatalf("identity not updated in state\nexpected: %s\ngot: %s", expectedIdentity, res.Current.IdentityJSON)
+	}
+	expectedVersion := uint64(5)
+	if res.Current.IdentitySchemaVersion != expectedVersion {
+		t.Fatalf("identity schema version not updated in state\nexpected: %d\ngot: %d", expectedVersion, res.Current.IdentitySchemaVersion)
+	}
+}
diff --git a/internal/terraform/upgrade_resource_state.go b/internal/terraform/upgrade_resource_state.go
index 7e6b698718..5529fc5574 100644
--- a/internal/terraform/upgrade_resource_state.go
+++ b/internal/terraform/upgrade_resource_state.go
@@ -182,7 +182,7 @@ func upgradeResourceIdentity(addr addrs.AbsResourceInstance, provider providers.
 		// to all protobuf target languages so in practice we use int64
 		// on the wire. In future we will change all of our internal
 		// representations to int64 too.
-		Version:         int64(src.SchemaVersion),
+		Version:         int64(src.IdentitySchemaVersion),
 		RawIdentityJSON: src.IdentityJSON,
 	}