From 12bb0f7510135c0869d042b418b03c8ebe036cd9 Mon Sep 17 00:00:00 2001
From: Venelin <venelin@pulumi.com>
Date: Mon, 25 Nov 2024 12:32:58 +0000
Subject: [PATCH] address review

---
 pkg/pf/internal/plugin/provider_server.go     | 38 ++++++++++++++-----
 .../internal/plugin/provider_server_test.go   | 37 ++++--------------
 2 files changed, 37 insertions(+), 38 deletions(-)

diff --git a/pkg/pf/internal/plugin/provider_server.go b/pkg/pf/internal/plugin/provider_server.go
index 17d8df2ca..4254c3309 100644
--- a/pkg/pf/internal/plugin/provider_server.go
+++ b/pkg/pf/internal/plugin/provider_server.go
@@ -23,9 +23,9 @@ import (
 	pbempty "github.com/golang/protobuf/ptypes/empty"
 	"github.com/pulumi/pulumi/sdk/v3/go/common/resource"
 	"github.com/pulumi/pulumi/sdk/v3/go/common/resource/config"
-	"github.com/pulumi/pulumi/sdk/v3/go/common/resource/plugin"
 	pl "github.com/pulumi/pulumi/sdk/v3/go/common/resource/plugin"
 	"github.com/pulumi/pulumi/sdk/v3/go/common/tokens"
+	"github.com/pulumi/pulumi/sdk/v3/go/common/util/contract"
 	pulumirpc "github.com/pulumi/pulumi/sdk/v3/proto/go"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -76,7 +76,27 @@ func (p *providerServer) checkNYI(method string, err error) error {
 	return err
 }
 
-func (p *providerServer) marshalDiff(diff plugin.DiffResult) (*pulumirpc.DiffResponse, error) {
+func pluginDiffKindToRPC(kind pl.DiffKind) pulumirpc.PropertyDiff_Kind {
+	switch kind {
+	case pl.DiffAdd:
+		return pulumirpc.PropertyDiff_ADD
+	case pl.DiffAddReplace:
+		return pulumirpc.PropertyDiff_ADD_REPLACE
+	case pl.DiffDelete:
+		return pulumirpc.PropertyDiff_DELETE
+	case pl.DiffDeleteReplace:
+		return pulumirpc.PropertyDiff_DELETE_REPLACE
+	case pl.DiffUpdate:
+		return pulumirpc.PropertyDiff_UPDATE
+	case pl.DiffUpdateReplace:
+		return pulumirpc.PropertyDiff_UPDATE_REPLACE
+	default:
+		contract.Assertf(false, "unknown diff kind: %v", kind)
+		return pulumirpc.PropertyDiff_ADD
+	}
+}
+
+func (p *providerServer) marshalDiff(diff pl.DiffResult) (*pulumirpc.DiffResponse, error) {
 	var changes pulumirpc.DiffResponse_DiffChanges
 	switch diff.Changes {
 	case pl.DiffNone:
@@ -94,7 +114,7 @@ func (p *providerServer) marshalDiff(diff plugin.DiffResult) (*pulumirpc.DiffRes
 		detailedDiff = make(map[string]*pulumirpc.PropertyDiff)
 		for path, diff := range diff.DetailedDiff {
 			detailedDiff[path] = &pulumirpc.PropertyDiff{
-				Kind:      pulumirpc.PropertyDiff_Kind(diff.Kind), //nolint:gosec
+				Kind:      pluginDiffKindToRPC(diff.Kind),
 				InputDiff: diff.InputDiff,
 			}
 		}
@@ -120,21 +140,21 @@ func (p *providerServer) marshalDiff(diff plugin.DiffResult) (*pulumirpc.DiffRes
 }
 
 type forwardServer struct {
-	plugin.UnimplementedProvider
+	pl.UnimplementedProvider
 
-	parameterize func(context.Context, plugin.ParameterizeRequest) (plugin.ParameterizeResponse, error)
+	parameterize func(context.Context, pl.ParameterizeRequest) (pl.ParameterizeResponse, error)
 }
 
 func (p forwardServer) Parameterize(
-	ctx context.Context, req plugin.ParameterizeRequest,
-) (plugin.ParameterizeResponse, error) {
+	ctx context.Context, req pl.ParameterizeRequest,
+) (pl.ParameterizeResponse, error) {
 	return p.parameterize(ctx, req)
 }
 
 func (p *providerServer) Parameterize(
 	ctx context.Context, req *pulumirpc.ParameterizeRequest,
 ) (*pulumirpc.ParameterizeResponse, error) {
-	return plugin.NewProviderServer(&forwardServer{
+	return pl.NewProviderServer(&forwardServer{
 		parameterize: p.provider.ParameterizeWithContext,
 	}).Parameterize(ctx, req)
 }
@@ -150,7 +170,7 @@ func (p *providerServer) GetSchema(ctx context.Context,
 		}
 		subpackageVersion = &ver
 	}
-	schema, err := p.provider.GetSchemaWithContext(ctx, plugin.GetSchemaRequest{
+	schema, err := p.provider.GetSchemaWithContext(ctx, pl.GetSchemaRequest{
 		Version:           req.GetVersion(),
 		SubpackageName:    req.SubpackageName,
 		SubpackageVersion: subpackageVersion,
diff --git a/pkg/pf/internal/plugin/provider_server_test.go b/pkg/pf/internal/plugin/provider_server_test.go
index c08070f42..dcac92f09 100644
--- a/pkg/pf/internal/plugin/provider_server_test.go
+++ b/pkg/pf/internal/plugin/provider_server_test.go
@@ -3,33 +3,12 @@ package plugin
 import (
 	"testing"
 
-	"github.com/hexops/autogold/v2"
 	"github.com/pulumi/pulumi/sdk/v3/go/common/resource"
 	"github.com/pulumi/pulumi/sdk/v3/go/common/resource/plugin"
 	pulumirpc "github.com/pulumi/pulumi/sdk/v3/proto/go"
 	"github.com/stretchr/testify/require"
 )
 
-// autogold gets confused about the enums, so we just re-define them here.
-//
-//nolint:revive
-const (
-	ADD            pulumirpc.PropertyDiff_Kind = 0
-	ADD_REPLACE    pulumirpc.PropertyDiff_Kind = 1
-	DELETE         pulumirpc.PropertyDiff_Kind = 2
-	DELETE_REPLACE pulumirpc.PropertyDiff_Kind = 3
-	UPDATE         pulumirpc.PropertyDiff_Kind = 4
-	UPDATE_REPLACE pulumirpc.PropertyDiff_Kind = 5
-)
-
-//nolint:revive
-const (
-	DIFF_UNKNOWN pulumirpc.DiffResponse_DiffChanges = 0
-	DIFF_NONE    pulumirpc.DiffResponse_DiffChanges = 1
-	DIFF_SOME    pulumirpc.DiffResponse_DiffChanges = 2
-)
-
-//nolint:unconvert
 func TestMarshalDiff(t *testing.T) {
 	t.Parallel()
 
@@ -48,11 +27,11 @@ func TestMarshalDiff(t *testing.T) {
 			DetailedDiff: map[string]plugin.PropertyDiff{},
 		}
 
-		autogold.Expect(&pulumirpc.DiffResponse{
+		require.Equal(t, &pulumirpc.DiffResponse{
 			Replaces: []string{},
 			Changes:  pulumirpc.DiffResponse_DIFF_NONE,
 			Diffs:    []string{},
-		}).Equal(t, runTest(t, diff))
+		}, runTest(t, diff))
 	})
 
 	t.Run("diff without detailed diff", func(t *testing.T) {
@@ -62,11 +41,11 @@ func TestMarshalDiff(t *testing.T) {
 			ChangedKeys: []resource.PropertyKey{"change"},
 		}
 
-		autogold.Expect(&pulumirpc.DiffResponse{
+		require.Equal(t, &pulumirpc.DiffResponse{
 			Replaces: []string{"replace"},
 			Changes:  pulumirpc.DiffResponse_DIFF_SOME,
 			Diffs:    []string{"change"},
-		}).Equal(t, runTest(t, diff))
+		}, runTest(t, diff))
 	})
 
 	t.Run("diff with detailed diff", func(t *testing.T) {
@@ -84,20 +63,20 @@ func TestMarshalDiff(t *testing.T) {
 			},
 		}
 
-		autogold.Expect(&pulumirpc.DiffResponse{
+		require.Equal(t, &pulumirpc.DiffResponse{
 			Replaces: []string{
 				"replace",
 			},
-			Changes: pulumirpc.DiffResponse_DiffChanges(DIFF_SOME),
+			Changes: pulumirpc.DiffResponse_DIFF_SOME,
 			Diffs: []string{
 				"change",
 				"replace",
 			},
 			DetailedDiff: map[string]*pulumirpc.PropertyDiff{
 				"change":  {},
-				"replace": {Kind: pulumirpc.PropertyDiff_Kind(DELETE_REPLACE)},
+				"replace": {Kind: pulumirpc.PropertyDiff_DELETE_REPLACE},
 			},
 			HasDetailedDiff: true,
-		}).Equal(t, runTest(t, diff))
+		}, runTest(t, diff))
 	})
 }