From b626f148ee400490fa598760f3ecfd6e5f303e05 Mon Sep 17 00:00:00 2001
From: chaitalisg <chaitali@kasten.io>
Date: Tue, 20 Sep 2022 14:00:53 -0700
Subject: [PATCH] test(general): Use new kopia upgrade cli in robustness
 upgrade tests (#2355)

* Use new kopia upgrade cli

* Clean up environement variable

* Clean up environement variable

* Clean up debugging change

* Fix lint errors

* Fix lint errors
---
 tests/robustness/robustness_test/main_test.go  |  5 +++++
 tests/robustness/snapmeta/kopia_snapshotter.go | 13 ++++++++++++-
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/tests/robustness/robustness_test/main_test.go b/tests/robustness/robustness_test/main_test.go
index 35bcc089d..320d809cf 100644
--- a/tests/robustness/robustness_test/main_test.go
+++ b/tests/robustness/robustness_test/main_test.go
@@ -70,6 +70,9 @@ func TestMain(m *testing.M) {
 
 		curr := rs.ContentFormat.MutableParameters.Version
 		log.Println("Upgraded repository format:", curr)
+
+		// Reset the env variable.
+		os.Setenv("KOPIA_UPGRADE_LOCK_ENABLED", "")
 	}
 
 	// run the tests
@@ -216,6 +219,8 @@ func (th *kopiaRobustnessTestHarness) getEngine() bool {
 }
 
 func (th *kopiaRobustnessTestHarness) cleanup(ctx context.Context) (retErr error) {
+	os.Setenv("KOPIA_UPGRADE_LOCK_ENABLED", "")
+
 	if th.engine != nil {
 		retErr = th.engine.Shutdown(ctx)
 	}
diff --git a/tests/robustness/snapmeta/kopia_snapshotter.go b/tests/robustness/snapmeta/kopia_snapshotter.go
index f84443298..1ea2a290d 100644
--- a/tests/robustness/snapmeta/kopia_snapshotter.go
+++ b/tests/robustness/snapmeta/kopia_snapshotter.go
@@ -8,6 +8,7 @@
 	"encoding/json"
 	"io"
 	"log"
+	"os"
 	"os/exec"
 	"strconv"
 
@@ -200,7 +201,17 @@ func (ks *KopiaSnapshotter) GetRepositoryStatus() (cli.RepositoryStatus, error)
 // UpgradeRepository upgrades the given kopia repository
 // from current format version to latest stable format version.
 func (ks *KopiaSnapshotter) UpgradeRepository() error {
-	_, _, err := ks.snap.Run("repository", "set-parameters", "--upgrade")
+	// This variable is also reset in cleanup function
+	// in case the test fails
+	os.Setenv("KOPIA_UPGRADE_LOCK_ENABLED", "1")
+
+	_, _, err := ks.snap.Run("repository", "upgrade",
+		"--upgrade-owner-id", "robustness-tests",
+		"--io-drain-timeout", "30s", "--allow-unsafe-upgrade",
+		"--status-poll-interval", "1s")
+
+	// cleanup
+	os.Setenv("KOPIA_UPGRADE_LOCK_ENABLED", "")
 
 	return err
 }