From af38e411496b323df68b316715f06172ab45c0ad Mon Sep 17 00:00:00 2001
From: Callum Waters <cmwaters19@gmail.com>
Date: Tue, 24 Sep 2024 18:39:38 +0200
Subject: [PATCH] feat: add cli command for using upgrade sequences with txsim
 (#3903)

We need this so our e2e tests can schedule an upgrade to a new network
version
---
 test/cmd/txsim/cli.go      | 35 +++++++++++++++++++++++++++++++++++
 test/cmd/txsim/cli_test.go |  1 +
 2 files changed, 36 insertions(+)

diff --git a/test/cmd/txsim/cli.go b/test/cmd/txsim/cli.go
index 62988c3eae..bf334437c4 100644
--- a/test/cmd/txsim/cli.go
+++ b/test/cmd/txsim/cli.go
@@ -40,6 +40,7 @@ var (
 	send, sendIterations, sendAmount                  int
 	stake, stakeValue, blob                           int
 	useFeegrant, suppressLogs                         bool
+	upgradeSchedule                                   string
 )
 
 func main() {
@@ -130,6 +131,15 @@ well funded account that can act as the master account. The command runs until a
 				sequences = append(sequences, txsim.NewBlobSequence(sizes, blobsPerPFB).Clone(blob)...)
 			}
 
+			upgradeScheduleMap, err := parseUpgradeSchedule(upgradeSchedule)
+			if err != nil {
+				return fmt.Errorf("invalid upgrade schedule: %w", err)
+			}
+
+			for height, version := range upgradeScheduleMap {
+				sequences = append(sequences, txsim.NewUpgradeSequence(version, height))
+			}
+
 			if seed == 0 {
 				if os.Getenv(TxsimSeed) != "" {
 					seed, err = strconv.ParseInt(os.Getenv(TxsimSeed), 10, 64)
@@ -195,6 +205,7 @@ func flags() *flag.FlagSet {
 	flags.IntVar(&stake, "stake", 0, "number of stake sequences to run")
 	flags.IntVar(&stakeValue, "stake-value", 1000, "amount of initial stake per sequence")
 	flags.IntVar(&blob, "blob", 0, "number of blob sequences to run")
+	flags.StringVar(&upgradeSchedule, "upgrade-schedule", "", "upgrade schedule for the network in format height:version i.e. 100:3,200:4")
 	flags.StringVar(&blobSizes, "blob-sizes", "100-1000", "range of blob sizes to send")
 	flags.StringVar(&blobAmounts, "blob-amounts", "1", "range of blobs per PFB specified as a single value or a min-max range (e.g., 10 or 5-10). A single value indicates the exact number of blobs to be created.")
 	flags.BoolVar(&useFeegrant, "feegrant", false, "use the feegrant module to pay for fees")
@@ -224,3 +235,27 @@ func readRange(r string) (txsim.Range, error) {
 
 	return txsim.NewRange(n, m), nil
 }
+
+func parseUpgradeSchedule(schedule string) (map[int64]uint64, error) {
+	scheduleMap := make(map[int64]uint64)
+	if schedule == "" {
+		return nil, nil
+	}
+	scheduleParts := strings.Split(schedule, ",")
+	for _, part := range scheduleParts {
+		parts := strings.Split(part, ":")
+		if len(parts) != 2 {
+			return nil, fmt.Errorf("invalid upgrade schedule format: %s", part)
+		}
+		height, err := strconv.ParseInt(parts[0], 10, 64)
+		if err != nil {
+			return nil, fmt.Errorf("invalid height in upgrade schedule: %s", parts[0])
+		}
+		version, err := strconv.ParseUint(parts[1], 10, 64)
+		if err != nil {
+			return nil, fmt.Errorf("invalid version in upgrade schedule: %s", parts[1])
+		}
+		scheduleMap[height] = version
+	}
+	return scheduleMap, nil
+}
diff --git a/test/cmd/txsim/cli_test.go b/test/cmd/txsim/cli_test.go
index 3520221af0..ad6d1c2268 100644
--- a/test/cmd/txsim/cli_test.go
+++ b/test/cmd/txsim/cli_test.go
@@ -29,6 +29,7 @@ func TestTxsimCommandFlags(t *testing.T) {
 		"--grpc-endpoint", grpcAddr,
 		"--blob", "5",
 		"--seed", "1234",
+		"--upgrade-schedule", "10:3",
 	})
 	err := cmd.ExecuteContext(ctx)
 	require.NoError(t, err)