Diff cross tests for provider upgrades (#2910)

This PR adds a `DiffProviderUpgradedSchema` option to `Diff` cross-tests
which allows us to test transitions between two different resource
schemas.

It uses the same `pulumitest` mechanism (`t.Run`) as the upgrade tests.

I've also added a repro for
#1667

Part of #2914

Author: VenelinMartinov

pkg/internal/tests/cross-tests/diff.go

@@ -10,6 +10,7 @@ import (
 type diffOpts struct {
 	deleteBeforeReplace           bool
 	disableAccurateBridgePreviews bool
+	resource2                     *schema.Resource
 }
 
 // An option that can be used to customize [Diff].
@@ -25,6 +26,11 @@ func DiffDisableAccurateBridgePreviews() DiffOption {
 	return func(o *diffOpts) { o.disableAccurateBridgePreviews = true }
 }
 
+// DiffProviderUpgradedSchema specifies the second provider schema to use for the diff.
+func DiffProviderUpgradedSchema(resource2 *schema.Resource) DiffOption {
+	return func(o *diffOpts) { o.resource2 = resource2 }
+}
+
 func Diff(
 	t T, resource *schema.Resource, config1, config2 map[string]cty.Value, opts ...DiffOption,
 ) crosstestsimpl.DiffResult {
@@ -42,5 +48,6 @@ func Diff(
 		Config2:                       config2Cty,
 		DeleteBeforeReplace:           o.deleteBeforeReplace,
 		DisableAccurateBridgePreviews: o.disableAccurateBridgePreviews,
+		Resource2:                     o.resource2,
 	})
 }

pkg/internal/tests/cross-tests/diff_check.go

@@ -15,11 +15,14 @@
 package crosstests
 
 import (
-	"os"
-	"path/filepath"
+	"context"
 
 	"github.com/hashicorp/terraform-plugin-go/tftypes"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
+	"github.com/pulumi/providertest/providers"
+	"github.com/pulumi/providertest/pulumitest"
+	"github.com/pulumi/providertest/pulumitest/optrun"
+	"github.com/pulumi/providertest/pulumitest/opttest"
 	"github.com/pulumi/pulumi/sdk/v3/go/auto/optpreview"
 	"github.com/stretchr/testify/require"
 
@@ -44,49 +47,76 @@ type diffTestCase struct {
 	ObjectType                    *tftypes.Object
 	DeleteBeforeReplace           bool
 	DisableAccurateBridgePreviews bool
+
+	// Optional second schema to use as an upgrade test with a different schema.
+	Resource2 *schema.Resource
 }
 
 func runDiffCheck(t T, tc diffTestCase) crosstestsimpl.DiffResult {
 	t.Helper()
 	tfwd := t.TempDir()
 
 	lifecycleArgs := lifecycleArgs{CreateBeforeDestroy: !tc.DeleteBeforeReplace}
+	resource1 := tc.Resource
+	resource2 := tc.Resource2
+	if resource2 == nil {
+		resource2 = resource1
+	}
+
+	tfConfig1 := coalesceInputs(t, resource1.Schema, tc.Config1)
+	tfd := newTFResDriver(t, tfwd, defProviderShortName, defRtype, resource1)
+	_ = tfd.writePlanApply(t, resource1.Schema, defRtype, "example", tfConfig1, lifecycleArgs)
 
-	tfConfig1 := coalesceInputs(t, tc.Resource.Schema, tc.Config1)
-	tfConfig2 := coalesceInputs(t, tc.Resource.Schema, tc.Config2)
-	tfd := newTFResDriver(t, tfwd, defProviderShortName, defRtype, tc.Resource)
-	_ = tfd.writePlanApply(t, tc.Resource.Schema, defRtype, "example", tfConfig1, lifecycleArgs)
-	tfDiffPlan := tfd.writePlanApply(t, tc.Resource.Schema, defRtype, "example", tfConfig2, lifecycleArgs)
+	tfConfig2 := coalesceInputs(t, resource2.Schema, tc.Config2)
+	tfd2 := newTFResDriver(t, tfwd, defProviderShortName, defRtype, resource2)
+	tfDiffPlan := tfd2.writePlanApply(t, resource2.Schema, defRtype, "example", tfConfig2, lifecycleArgs)
 
-	resMap := map[string]*schema.Resource{defRtype: tc.Resource}
-	tfp := &schema.Provider{ResourcesMap: resMap}
+	tfp1 := &schema.Provider{ResourcesMap: map[string]*schema.Resource{defRtype: resource1}}
+	tfp2 := &schema.Provider{ResourcesMap: map[string]*schema.Resource{defRtype: resource2}}
 
 	opts := []pulcheck.BridgedProviderOpt{}
 	if !tc.DisableAccurateBridgePreviews {
 		opts = append(opts, pulcheck.EnableAccurateBridgePreviews())
 	}
 
-	bridgedProvider := pulcheck.BridgedProvider(t, defProviderShortName, tfp, opts...)
+	bridgedProvider1 := pulcheck.BridgedProvider(t, defProviderShortName, tfp1, opts...)
+	bridgedProvider2 := pulcheck.BridgedProvider(t, defProviderShortName, tfp2, opts...)
 	if tc.DeleteBeforeReplace {
-		bridgedProvider.Resources[defRtype].DeleteBeforeReplace = true
+		bridgedProvider1.Resources[defRtype].DeleteBeforeReplace = true
+		bridgedProvider2.Resources[defRtype].DeleteBeforeReplace = true
 	}
 
 	pd := &pulumiDriver{
 		name:                defProviderShortName,
 		pulumiResourceToken: defRtoken,
 		tfResourceName:      defRtype,
 	}
-
-	yamlProgram := pd.generateYAML(t, crosstestsimpl.InferPulumiValue(t,
-		bridgedProvider.P.ResourcesMap().Get(defRtype).Schema(), nil, tfConfig1))
-	pt := pulcheck.PulCheck(t, bridgedProvider, string(yamlProgram))
-	pt.Up(t)
-
-	yamlProgram = pd.generateYAML(t, crosstestsimpl.InferPulumiValue(t,
-		bridgedProvider.P.ResourcesMap().Get(defRtype).Schema(), nil, tfConfig2))
-	err := os.WriteFile(filepath.Join(pt.CurrentStack().Workspace().WorkDir(), "Pulumi.yaml"), yamlProgram, 0o600)
-	require.NoErrorf(t, err, "writing Pulumi.yaml")
-
+	yamlProgram1 := pd.generateYAML(t, crosstestsimpl.InferPulumiValue(t,
+		bridgedProvider1.P.ResourcesMap().Get(defRtype).Schema(), nil, tfConfig1))
+
+	yamlProgram2 := pd.generateYAML(t, crosstestsimpl.InferPulumiValue(t,
+		bridgedProvider2.P.ResourcesMap().Get(defRtype).Schema(), nil, tfConfig2))
+
+	// We initialize the second provider as it will be used in the preview.
+	// It is temporarily overwritten by the first provider in the Run function.
+	pt := pulcheck.PulCheck(t, bridgedProvider2, string(yamlProgram1))
+	pt.Run(
+		t,
+		func(test *pulumitest.PulumiTest) {
+			test.Up(t)
+		},
+		optrun.WithOpts(
+			opttest.AttachProvider(
+				defProviderShortName,
+				func(ctx context.Context, pt providers.PulumiTest) (providers.Port, error) {
+					handle, err := pulcheck.StartPulumiProvider(ctx, bridgedProvider1)
+					require.NoError(t, err)
+					return providers.Port(handle.Port), nil
+				},
+			),
+		),
+	)
+	pt.WritePulumiYaml(t, string(yamlProgram2))
 	previewRes := pt.Preview(t, optpreview.Diff())
 	require.Empty(t, previewRes.StdErr, "preview should not have errors")
 

pkg/internal/tests/cross-tests/diff_cross_test.go

@@ -20,6 +20,7 @@ import (
 	"fmt"
 	"hash/crc32"
 	"slices"
+	"strconv"
 	"strings"
 	"testing"
 
@@ -1734,3 +1735,55 @@ func TestDetailedDiffReplacementComputedProperty(t *testing.T) {
 		})
 	})
 }
+
+func TestDiffProviderUpgradeBasic(t *testing.T) {
+	t.Parallel()
+
+	res1 := &schema.Resource{
+		Schema: map[string]*schema.Schema{"prop": {Type: schema.TypeString, Optional: true}},
+	}
+
+	res2 := &schema.Resource{
+		Schema:        map[string]*schema.Schema{"prop": {Type: schema.TypeInt, Optional: true}},
+		SchemaVersion: 1,
+		StateUpgraders: []schema.StateUpgrader{
+			{
+				Version: 0,
+				Type:    res1.CoreConfigSchema().ImpliedType(),
+				Upgrade: func(ctx context.Context, rawState map[string]any, meta interface{}) (map[string]any, error) {
+					if rawState == nil {
+						rawState = map[string]interface{}{}
+					}
+
+					if _, ok := rawState["prop"]; ok {
+						stringVar, ok := rawState["prop"].(string)
+						var intVar int
+						// TODO[pulumi/pulumi-terraform-bridge#1667]: This is a workaround to handle
+						// the fact that we use the new schema to decode the state
+						if !ok {
+							floatVar := rawState["prop"].(float64)
+							intVar = int(floatVar)
+						} else {
+							var err error
+							intVar, err = strconv.Atoi(stringVar)
+							if err != nil {
+								return nil, err
+							}
+						}
+						rawState["prop"] = intVar
+					}
+
+					return rawState, nil
+				},
+			},
+		},
+	}
+
+	res := Diff(t, res1,
+		map[string]cty.Value{"prop": cty.StringVal("1")},
+		map[string]cty.Value{"prop": cty.NumberIntVal(1)},
+		DiffProviderUpgradedSchema(res2),
+	)
+
+	require.Equal(t, []string{"no-op"}, res.TFDiff.Actions)
+}