diff --git a/internal/terraform/context_plan_identity_test.go b/internal/terraform/context_plan_identity_test.go
index 64a135db5b19..cafc3531fe33 100644
--- a/internal/terraform/context_plan_identity_test.go
+++ b/internal/terraform/context_plan_identity_test.go
@@ -216,6 +216,29 @@ func TestContext2Plan_resource_identity_refresh(t *testing.T) {
 		t.Run(name, func(t *testing.T) {
 			p := testProvider("aws")
 			m := testModule(t, "refresh-basic")
+			p.GetProviderSchemaResponse = getProviderSchemaResponseFromProviderSchema(&providerSchema{
+				ResourceTypes: map[string]*configschema.Block{
+					"aws_instance": {
+						Attributes: map[string]*configschema.Attribute{
+							"id": {
+								Type:     cty.String,
+								Computed: true,
+							},
+							"foo": {
+								Type:     cty.String,
+								Optional: true,
+								Computed: true,
+							},
+						},
+					},
+				},
+				IdentityTypes: map[string]*configschema.Object{
+					"aws_instance": tc.IdentitySchema.Body,
+				},
+				IdentityTypeSchemaVersions: map[string]uint64{
+					"aws_instance": uint64(tc.IdentitySchema.Version),
+				},
+			})
 
 			state := states.NewState()
 			root := state.EnsureModule(addrs.RootModuleInstance)
@@ -244,12 +267,6 @@ func TestContext2Plan_resource_identity_refresh(t *testing.T) {
 				t.Fatal(err)
 			}
 
-			p.GetResourceIdentitySchemasResponse = &providers.GetResourceIdentitySchemasResponse{
-				IdentityTypes: map[string]providers.IdentitySchema{
-					"aws_instance": tc.IdentitySchema,
-				},
-			}
-			schema.Identity = p.GetResourceIdentitySchemasResponse.IdentityTypes["aws_instance"].Body
 			p.ReadResourceResponse = &providers.ReadResourceResponse{
 				NewState: readState,
 				Identity: tc.IdentityData,
@@ -279,10 +296,6 @@ func TestContext2Plan_resource_identity_refresh(t *testing.T) {
 				t.Fatal("ReadResource should be called")
 			}
 
-			if !p.GetResourceIdentitySchemasCalled {
-				t.Fatal("GetResourceIdentitySchemas should be called")
-			}
-
 			if tc.ExpectUpgradeResourceIdentityCalled && !p.UpgradeResourceIdentityCalled {
 				t.Fatal("UpgradeResourceIdentity should be called")
 			}
@@ -314,6 +327,37 @@ func TestContext2Plan_resource_identity_refresh(t *testing.T) {
 func TestContext2Plan_resource_identity_refresh_destroy_deposed(t *testing.T) {
 	p := testProvider("aws")
 	m := testModule(t, "refresh-basic")
+	p.GetProviderSchemaResponse = getProviderSchemaResponseFromProviderSchema(&providerSchema{
+		ResourceTypes: map[string]*configschema.Block{
+			"aws_instance": {
+				Attributes: map[string]*configschema.Attribute{
+					"id": {
+						Type:     cty.String,
+						Computed: true,
+					},
+					"foo": {
+						Type:     cty.String,
+						Optional: true,
+						Computed: true,
+					},
+				},
+			},
+		},
+		IdentityTypes: map[string]*configschema.Object{
+			"aws_instance": {
+				Attributes: map[string]*configschema.Attribute{
+					"id": {
+						Type:     cty.String,
+						Required: true,
+					},
+				},
+				Nesting: configschema.NestingSingle,
+			},
+		},
+		IdentityTypeSchemaVersions: map[string]uint64{
+			"aws_instance": 0,
+		},
+	})
 
 	state := states.NewState()
 	root := state.EnsureModule(addrs.RootModuleInstance)
@@ -342,23 +386,6 @@ func TestContext2Plan_resource_identity_refresh_destroy_deposed(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	p.GetResourceIdentitySchemasResponse = &providers.GetResourceIdentitySchemasResponse{
-		IdentityTypes: map[string]providers.IdentitySchema{
-			"aws_instance": {
-				Version: 0,
-				Body: &configschema.Object{
-					Attributes: map[string]*configschema.Attribute{
-						"id": {
-							Type:     cty.String,
-							Required: true,
-						},
-					},
-					Nesting: configschema.NestingSingle,
-				},
-			},
-		},
-	}
-	schema.Identity = p.GetResourceIdentitySchemasResponse.IdentityTypes["aws_instance"].Body
 	p.ReadResourceResponse = &providers.ReadResourceResponse{
 		NewState: readState,
 		Identity: cty.ObjectVal(map[string]cty.Value{
@@ -376,10 +403,6 @@ func TestContext2Plan_resource_identity_refresh_destroy_deposed(t *testing.T) {
 		t.Fatal("ReadResource should be called")
 	}
 
-	if !p.GetResourceIdentitySchemasCalled {
-		t.Fatal("GetResourceIdentitySchemas should be called")
-	}
-
 	mod := s.PriorState.RootModule()
 	fromState, err := mod.Resources["aws_instance.web"].Instances[addrs.NoKey].Deposed[deposedKey].Decode(schema)
 	if err != nil {
diff --git a/internal/terraform/resource_provider_mock_test.go b/internal/terraform/resource_provider_mock_test.go
index f0b8006bd8a4..45da48cbf1b5 100644
--- a/internal/terraform/resource_provider_mock_test.go
+++ b/internal/terraform/resource_provider_mock_test.go
@@ -124,6 +124,8 @@ type providerSchema struct {
 	ResourceTypes              map[string]*configschema.Block
 	ResourceTypeSchemaVersions map[string]uint64
 	DataSources                map[string]*configschema.Block
+	IdentityTypes              map[string]*configschema.Object
+	IdentityTypeSchemaVersions map[string]uint64
 }
 
 // getProviderSchemaResponseFromProviderSchema is a test helper to convert a
@@ -137,10 +139,18 @@ func getProviderSchemaResponseFromProviderSchema(providerSchema *providerSchema)
 	}
 
 	for name, schema := range providerSchema.ResourceTypes {
-		resp.ResourceTypes[name] = providers.Schema{
+		ps := providers.Schema{
 			Body:    schema,
 			Version: int64(providerSchema.ResourceTypeSchemaVersions[name]),
 		}
+
+		id, ok := providerSchema.IdentityTypes[name]
+		if ok {
+			ps.Identity = id
+			ps.IdentityVersion = int64(providerSchema.IdentityTypeSchemaVersions[name])
+		}
+
+		resp.ResourceTypes[name] = ps
 	}
 
 	for name, schema := range providerSchema.DataSources {