mock GetResourceIdentitySchemas

hashicorp · Mar 5, 2025 · 0f0c535 · 0f0c535
1 parent 621267b
commit 0f0c535
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 0 deletions.
diff --git a/internal/plugin/grpc_provider_test.go b/internal/plugin/grpc_provider_test.go
@@ -35,6 +35,13 @@ func mockProviderClient(t *testing.T) *mockproto.MockProviderClient {
 		gomock.Any(),
 	).Return(providerProtoSchema(), nil)
 
+	// GetResourceIdentitySchemas is called as part of GetSchema
+	client.EXPECT().GetResourceIdentitySchemas(
+		gomock.Any(),
+		gomock.Any(),
+		gomock.Any(),
+	).Return(providerResourceIdentitySchemas(), nil)
+
 	return client
 }
 
@@ -114,6 +121,22 @@ func providerProtoSchema() *proto.GetProviderSchema_Response {
 	}
 }
 
+func providerResourceIdentitySchemas() *proto.GetResourceIdentitySchemas_Response {
+	return &proto.GetResourceIdentitySchemas_Response{
+		IdentitySchemas: map[string]*proto.ResourceIdentitySchema{
+			"resource": {
+				Version: 1,
+				IdentityAttributes: []*proto.ResourceIdentitySchema_IdentityAttribute{
+					{
+						Name:              "attr",
+						RequiredForImport: true,
+					},
+				},
+			},
+		},
+	}
+}
+
 func TestGRPCProvider_GetSchema(t *testing.T) {
 	p := &GRPCProvider{
 		client: mockProviderClient(t),

diff --git a/internal/plugin6/grpc_provider_test.go b/internal/plugin6/grpc_provider_test.go
@@ -42,6 +42,13 @@ func mockProviderClient(t *testing.T) *mockproto.MockProviderClient {
 		gomock.Any(),
 	).Return(providerProtoSchema(), nil)
 
+	// GetResourceIdentitySchemas is called as part of GetSchema
+	client.EXPECT().GetResourceIdentitySchemas(
+		gomock.Any(),
+		gomock.Any(),
+		gomock.Any(),
+	).Return(providerResourceIdentitySchemas(), nil)
+
 	return client
 }
 
@@ -121,6 +128,22 @@ func providerProtoSchema() *proto.GetProviderSchema_Response {
 	}
 }
 
+func providerResourceIdentitySchemas() *proto.GetResourceIdentitySchemas_Response {
+	return &proto.GetResourceIdentitySchemas_Response{
+		IdentitySchemas: map[string]*proto.ResourceIdentitySchema{
+			"resource": {
+				Version: 1,
+				IdentityAttributes: []*proto.ResourceIdentitySchema_IdentityAttribute{
+					{
+						Name:              "attr",
+						RequiredForImport: true,
+					},
+				},
+			},
+		},
+	}
+}
+
 func TestGRPCProvider_GetSchema(t *testing.T) {
 	p := &GRPCProvider{
 		client: mockProviderClient(t),