From 418ddbfe4bcde2836e3460caba9add61f97a5709 Mon Sep 17 00:00:00 2001
From: Callum Waters <cmwaters19@gmail.com>
Date: Tue, 24 Sep 2024 11:37:50 +0200
Subject: [PATCH] feat: add upgrade sequence to tx sim (#3890)

Adds the ability to use txsim to upgrade a network
---
 test/txsim/run_test.go | 52 +++++++++++++++++++++++++++
 test/txsim/upgrade.go  | 82 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 134 insertions(+)
 create mode 100644 test/txsim/upgrade.go

diff --git a/test/txsim/run_test.go b/test/txsim/run_test.go
index 46e4523710..4a922f15a1 100644
--- a/test/txsim/run_test.go
+++ b/test/txsim/run_test.go
@@ -12,12 +12,17 @@ import (
 
 	"github.com/celestiaorg/celestia-app/v3/app"
 	"github.com/celestiaorg/celestia-app/v3/app/encoding"
+	v2 "github.com/celestiaorg/celestia-app/v3/pkg/appconsts/v2"
+	v3 "github.com/celestiaorg/celestia-app/v3/pkg/appconsts/v3"
 	"github.com/celestiaorg/celestia-app/v3/test/txsim"
 	"github.com/celestiaorg/celestia-app/v3/test/util/testnode"
 	"github.com/cosmos/cosmos-sdk/crypto/keyring"
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
 
 	blob "github.com/celestiaorg/celestia-app/v3/x/blob/types"
+	signaltypes "github.com/celestiaorg/celestia-app/v3/x/signal/types"
 	bank "github.com/cosmos/cosmos-sdk/x/bank/types"
 	distribution "github.com/cosmos/cosmos-sdk/x/distribution/types"
 	staking "github.com/cosmos/cosmos-sdk/x/staking/types"
@@ -152,3 +157,50 @@ func Setup(t testing.TB) (keyring.Keyring, string, string) {
 
 	return cctx.Keyring, rpcAddr, grpcAddr
 }
+
+func TestTxSimUpgrade(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping TestTxSimUpgrade in short mode.")
+	}
+	cp := app.DefaultConsensusParams()
+	cp.Version.AppVersion = v2.Version
+	cfg := testnode.DefaultConfig().
+		WithTimeoutCommit(300 * time.Millisecond).
+		WithConsensusParams(cp).
+		WithFundedAccounts("txsim-master")
+	cctx, _, grpcAddr := testnode.NewNetwork(t, cfg)
+
+	// updrade to v3 at height 20
+	sequences := []txsim.Sequence{
+		txsim.NewUpgradeSequence(v3.Version, 20),
+	}
+
+	opts := txsim.DefaultOptions().
+		// SuppressLogs().
+		WithPollTime(time.Millisecond * 100)
+
+	err := txsim.Run(
+		cctx.GoContext(),
+		grpcAddr,
+		cctx.Keyring,
+		encoding.MakeConfig(app.ModuleEncodingRegisters...),
+		opts,
+		sequences...,
+	)
+	require.NoError(t, err)
+
+	conn, err := grpc.NewClient(grpcAddr, grpc.WithTransportCredentials(insecure.NewCredentials()))
+	require.NoError(t, err)
+	defer conn.Close()
+
+	querier := signaltypes.NewQueryClient(conn)
+
+	// We can't check that the upgrade was successful because the upgrade height is thousands of blocks away
+	// and even at 300 millisecond block times, it would take too long. Instead we just want to assert
+	// that the upgrade is ready to be performed
+	require.Eventually(t, func() bool {
+		upgradePlan, err := querier.GetUpgrade(cctx.GoContext(), &signaltypes.QueryGetUpgradeRequest{})
+		require.NoError(t, err)
+		return upgradePlan.Upgrade != nil && upgradePlan.Upgrade.AppVersion == v3.Version
+	}, time.Second*20, time.Millisecond*100)
+}
diff --git a/test/txsim/upgrade.go b/test/txsim/upgrade.go
new file mode 100644
index 0000000000..0811f32232
--- /dev/null
+++ b/test/txsim/upgrade.go
@@ -0,0 +1,82 @@
+package txsim
+
+import (
+	"context"
+	"errors"
+	"math/rand"
+
+	signaltypes "github.com/celestiaorg/celestia-app/v3/x/signal/types"
+	"github.com/cosmos/cosmos-sdk/types"
+	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
+	"github.com/gogo/protobuf/grpc"
+)
+
+var _ Sequence = &UpgradeSequence{}
+
+const fundsForUpgrade = 100_000
+
+// UpgradeSequence simulates an upgrade proposal and voting process
+type UpgradeSequence struct {
+	voted       map[string]bool
+	height      int64
+	version     uint64
+	account     types.AccAddress
+	hasUpgraded bool
+}
+
+func NewUpgradeSequence(version uint64, height int64) *UpgradeSequence {
+	return &UpgradeSequence{version: version, height: height, voted: make(map[string]bool)}
+}
+
+func (s *UpgradeSequence) Clone(_ int) []Sequence {
+	panic("cloning not supported for upgrade sequence. Only a single sequence is needed")
+}
+
+// this is a no-op for the upgrade sequence
+func (s *UpgradeSequence) Init(_ context.Context, _ grpc.ClientConn, allocateAccounts AccountAllocator, _ *rand.Rand, _ bool) {
+	s.account = allocateAccounts(1, fundsForUpgrade)[0]
+}
+
+func (s *UpgradeSequence) Next(ctx context.Context, querier grpc.ClientConn, _ *rand.Rand) (Operation, error) {
+	if s.hasUpgraded {
+		return Operation{}, ErrEndOfSequence
+	}
+
+	stakingQuerier := stakingtypes.NewQueryClient(querier)
+	validatorsResp, err := stakingQuerier.Validators(ctx, &stakingtypes.QueryValidatorsRequest{})
+	if err != nil {
+		return Operation{}, err
+	}
+
+	if len(validatorsResp.Validators) == 0 {
+		return Operation{}, errors.New("no validators found")
+	}
+
+	// Choose a random validator to be the authority
+	var msg types.Msg
+	for _, validator := range validatorsResp.Validators {
+		if !s.voted[validator.OperatorAddress] {
+			msg = &signaltypes.MsgSignalVersion{
+				ValidatorAddress: validator.OperatorAddress,
+				Version:          s.version,
+			}
+			s.voted[validator.OperatorAddress] = true
+		}
+	}
+	// if all validators have voted, we can now try to upgrade.
+	if msg == nil {
+		msg = signaltypes.NewMsgTryUpgrade(s.account)
+		s.hasUpgraded = true
+	}
+
+	delay := uint64(0)
+	// apply a delay to the first sequence only
+	if len(s.voted) == 0 {
+		delay = uint64(s.height)
+	}
+
+	return Operation{
+		Msgs:  []types.Msg{msg},
+		Delay: delay,
+	}, nil
+}