From a7e6897d592b0b3038f108e8a0fc0c1a5ad6b8e8 Mon Sep 17 00:00:00 2001
From: Venelin <venelin@pulumi.com>
Date: Thu, 1 Aug 2024 15:14:47 +0300
Subject: [PATCH] upgrade tests with known source edits

---
 optproviderupgrade/optproviderupgrade.go |  7 +++++++
 previewProviderUpgrade.go                |  4 ++++
 previewProviderUpgrade_test.go           | 15 +++++++++++++++
 3 files changed, 26 insertions(+)

diff --git a/optproviderupgrade/optproviderupgrade.go b/optproviderupgrade/optproviderupgrade.go
index 5ffb4c1..8e2777f 100644
--- a/optproviderupgrade/optproviderupgrade.go
+++ b/optproviderupgrade/optproviderupgrade.go
@@ -38,10 +38,17 @@ func CacheDir(elem ...string) PreviewProviderUpgradeOpt {
 	})
 }
 
+func NewSourcePath(path string) PreviewProviderUpgradeOpt {
+	return optionFunc(func(o *PreviewProviderUpgradeOptions) {
+		o.NewSourcePath = path
+	})
+}
+
 type PreviewProviderUpgradeOptions struct {
 	CacheDirTemplate []string
 	DisableAttach    bool
 	BaselineOpts     []opttest.Option
+	NewSourcePath    string
 }
 
 type PreviewProviderUpgradeOpt interface {
diff --git a/previewProviderUpgrade.go b/previewProviderUpgrade.go
index e53067d..80f19a4 100644
--- a/previewProviderUpgrade.go
+++ b/previewProviderUpgrade.go
@@ -39,6 +39,10 @@ func PreviewProviderUpgrade(t pulumitest.PT, pulumiTest *pulumitest.PulumiTest,
 			baselineProviderOpt(options, providerName, baselineVersion)),
 		optrun.WithOpts(options.BaselineOpts...),
 	)
+
+	if options.NewSourcePath != "" {
+		previewTest.UpdateSource(options.NewSourcePath)
+	}
 	return previewTest.Preview()
 }
 
diff --git a/previewProviderUpgrade_test.go b/previewProviderUpgrade_test.go
index 4eb92c2..6c89965 100644
--- a/previewProviderUpgrade_test.go
+++ b/previewProviderUpgrade_test.go
@@ -29,3 +29,18 @@ func TestPreviewUpgradeCached(t *testing.T) {
 		optproviderupgrade.DisableAttach())
 	assert.Equal(t, uncachedPreviewResult, cachedPreviewResult, "expected uncached and cached preview to be the same")
 }
+
+func TestPreviewUpgradeWithKnownSourceEdit(t *testing.T) {
+	t.Parallel()
+	cacheDir := t.TempDir()
+	test := pulumitest.NewPulumiTest(t, filepath.Join("pulumitest", "testdata", "yaml_program"),
+		opttest.DownloadProviderVersion("random", "4.15.0"))
+
+	previewResult := providertest.PreviewProviderUpgrade(t, test, "random", "4.5.0",
+		optproviderupgrade.CacheDir(cacheDir, "{programName}", "{baselineVersion}"),
+		optproviderupgrade.DisableAttach(),
+		optproviderupgrade.NewSourcePath(filepath.Join("pulumitest", "testdata", "yaml_program_updated")),
+	)
+
+	assert.Contains(t, previewResult.StdOut, "random:index:RandomPassword password create")
+}