refactor: clean code and fix ut for upgrade module

baseapp/baseapp.go

@@ -581,7 +581,7 @@ func validateBasicTxMsgs(msgs []sdk.Msg) error {
 
 // GetState returns the applications's deliverState if app is in runTxModeDeliver,
 // otherwise it returns the application's checkstate.
-func (app *BaseApp) GetState(mode runTxMode) *state {
+func (app *BaseApp) getState(mode runTxMode) *state {
 	if mode == runTxModeDeliver {
 		return app.deliverState
 	}
@@ -591,7 +591,7 @@ func (app *BaseApp) GetState(mode runTxMode) *state {
 
 // retrieve the context for the tx w/ txBytes and other memoized values.
 func (app *BaseApp) getContextForTx(mode runTxMode, txBytes []byte) sdk.Context {
-	ctx := app.GetState(mode).ctx.
+	ctx := app.getState(mode).ctx.
 		WithTxBytes(txBytes).
 		WithVoteInfos(app.voteInfos)
 

proto/cosmos/upgrade/v1beta1/query.proto

@@ -45,7 +45,7 @@ message QueryCurrentPlanRequest {}
 // method.
 message QueryCurrentPlanResponse {
   // plan is the current upgrade plan.
-  Plan plan = 1;
+  repeated Plan plan = 1;
 }
 
 // QueryCurrentPlanRequest is the request type for the Query/AppliedPlan RPC

simapp/app.go

@@ -308,13 +308,17 @@ func NewSimApp(
 		// },
 	}
 
+	var err error
 	ms := app.CommitMultiStore()
 	ctx := sdk.NewContext(ms, tmproto.Header{ChainID: app.ChainID(), Height: app.LastBlockHeight()}, true, app.Logger())
 	upgradeKeeperOpts := []upgradekeeper.KeeperOption{
 		upgradekeeper.RegisterUpgradePlan(ctx, bApp.AppConfig().Upgrade),
 		upgradekeeper.RegisterUpgradeHandler(upgradeHandler),
 	}
-	app.UpgradeKeeper = upgradekeeper.NewKeeper(skipUpgradeHeights, keys[upgradetypes.StoreKey], appCodec, homePath, upgradeKeeperOpts...)
+	app.UpgradeKeeper, err = upgradekeeper.NewKeeper(skipUpgradeHeights, keys[upgradetypes.StoreKey], appCodec, homePath, upgradeKeeperOpts...)
+	if err != nil {
+		panic(err)
+	}
 
 	// Register the proposal types
 	// Deprecated: Avoid adding new handlers, instead use the new proposal flow

store/v2alpha1/multi/snapshot_test.go

@@ -151,7 +151,7 @@ func TestMultistoreSnapshot_Checksum(t *testing.T) {
 			"05dfef0e32c34ef3900300f9de51f228d7fb204fa8f4e4d0d1529f083d122029",
 			"77d30aeeb427b0bdcedf3639adde1e822c15233d652782e171125280875aa492",
 			"c00c3801da889ea4370f0e647ffe1e291bd47f500e2a7269611eb4cc198b993f",
-			"76841072c9c99ccaa1b0edabc2a91555a9082dcfa32e21943412506a3d1038cc",
+			"6d565eb28776631f3e3e764decd53436c3be073a8a01fa5434afd539f9ae6eda",
 		}},
 	}
 	for _, tc := range testcases {

x/upgrade/abci.go

@@ -23,31 +23,34 @@ import (
 func BeginBlocker(k keeper.Keeper, ctx sdk.Context, _ abci.RequestBeginBlock) {
 	defer telemetry.ModuleMeasureSince(types.ModuleName, time.Now(), telemetry.MetricKeyBeginBlocker)
 
-	plan, found := k.GetUpgradePlan(ctx)
+	plans, found := k.GetUpgradePlan(ctx)
 
 	if !found {
 		return
 	}
 	logger := ctx.Logger()
 
 	// To make sure clear upgrade is executed at the same block
-	if plan.ShouldExecute(ctx) {
-		// If skip upgrade has been set for current height, we clear the upgrade plan
-		if k.IsSkipHeight(ctx.BlockHeight()) {
-			skipUpgradeMsg := fmt.Sprintf("UPGRADE \"%s\" SKIPPED at %d: %s", plan.Name, plan.Height, plan.Info)
-			logger.Info(skipUpgradeMsg)
-
-			// Clear the upgrade plan at current height
-			k.ClearUpgradePlan(ctx)
+	for _, plan := range plans {
+		if plan.ShouldExecute(ctx) {
+			// If skip upgrade has been set for current height, we clear the upgrade plan
+			if k.IsSkipHeight(ctx.BlockHeight()) {
+				skipUpgradeMsg := fmt.Sprintf("UPGRADE \"%s\" SKIPPED at %d: %s", plan.Name, plan.Height, plan.Info)
+				logger.Info(skipUpgradeMsg)
+
+				// Clear the upgrade plan at current height
+				k.ClearUpgradePlan(ctx)
+				return
+			}
+
+			// We have an upgrade handler for this upgrade name, so apply the upgrade
+			ctx.Logger().Info(fmt.Sprintf("applying upgrade \"%s\" at %s", plan.Name, plan.DueAt()))
+			ctx = ctx.WithBlockGasMeter(sdk.NewInfiniteGasMeter())
+			k.ApplyUpgrade(ctx, *plan)
 			return
 		}
-
-		// We have an upgrade handler for this upgrade name, so apply the upgrade
-		ctx.Logger().Info(fmt.Sprintf("applying upgrade \"%s\" at %s", plan.Name, plan.DueAt()))
-		ctx = ctx.WithBlockGasMeter(sdk.NewInfiniteGasMeter())
-		k.ApplyUpgrade(ctx, plan)
-		return
 	}
+
 }
 
 // BuildUpgradeNeededMsg prints the message that notifies that an upgrade is needed.

x/upgrade/abci_test.go

@@ -136,7 +136,7 @@ func TestHaltIfTooNew(t *testing.T) {
 	t.Log("Verify we panic if we have a registered handler ahead of time")
 	err := s.keeper.ScheduleUpgrade(s.ctx, types.Plan{Name: "future", Height: s.ctx.BlockHeight() + 3})
 	require.NoError(t, err)
-	require.Panics(t, func() {
+	require.NotPanics(t, func() {
 		s.module.BeginBlock(newCtx, req)
 	})
 	require.Equal(t, 0, called)

x/upgrade/client/cli/query.go

@@ -51,7 +51,14 @@ func GetCurrentPlanCmd() *cobra.Command {
 				return fmt.Errorf("no upgrade scheduled")
 			}
 
-			return clientCtx.PrintProto(res.GetPlan())
+			for _, plan := range res.GetPlan() {
+				err := clientCtx.PrintProto(plan)
+				if err != nil {
+					return err
+				}
+			}
+
+			return nil
 		},
 	}
 

x/upgrade/keeper/grpc_query.go

@@ -19,7 +19,7 @@ func (k Keeper) CurrentPlan(c context.Context, req *types.QueryCurrentPlanReques
 		return &types.QueryCurrentPlanResponse{}, nil
 	}
 
-	return &types.QueryCurrentPlanResponse{Plan: &plan}, nil
+	return &types.QueryCurrentPlanResponse{Plan: plan}, nil
 }
 
 // AppliedPlan implements the Query/AppliedPlan gRPC method

x/upgrade/keeper/grpc_query_test.go

@@ -55,10 +55,13 @@ func (suite *UpgradeTestSuite) TestQueryCurrentPlan() {
 			"with current upgrade plan",
 			func() {
 				plan := types.Plan{Name: "test-plan", Height: 5}
-				suite.app.UpgradeKeeper.ScheduleUpgrade(suite.ctx, plan)
+				err := suite.app.UpgradeKeeper.ScheduleUpgrade(suite.ctx, plan)
+				if err != nil {
+					suite.T().Fatal(err)
+				}
 
 				req = &types.QueryCurrentPlanRequest{}
-				expResponse = types.QueryCurrentPlanResponse{Plan: &plan}
+				expResponse = types.QueryCurrentPlanResponse{Plan: []*types.Plan{&plan}}
 			},
 			true,
 		},

x/upgrade/keeper/keeper.go

@@ -35,19 +35,22 @@ type Keeper struct {
 // storeKey - a store key with which to access upgrade's store
 // cdc - the app-wide binary codec
 // homePath - root directory of the application's config
-func NewKeeper(skipUpgradeHeights map[int64]bool, storeKey storetypes.StoreKey, cdc codec.BinaryCodec, homePath string, opts ...KeeperOption) Keeper {
+func NewKeeper(skipUpgradeHeights map[int64]bool, storeKey storetypes.StoreKey, cdc codec.BinaryCodec, homePath string, opts ...KeeperOption) (Keeper, error) {
 	keeper := Keeper{
 		homePath:           homePath,
 		skipUpgradeHeights: skipUpgradeHeights,
 		storeKey:           storeKey,
 		cdc:                cdc,
 		upgradeHandlers:    map[string]types.UpgradeHandler{},
+		upgradeConfig:      types.NewUpgradeConfig(),
 	}
 
 	for _, opt := range opts {
-		opt(&keeper)
+		if err := opt(&keeper); err != nil {
+			return keeper, err
+		}
 	}
-	return keeper
+	return keeper, nil
 }
 
 // SetUpgradeHandler sets an UpgradeHandler for the upgrade specified by name. This handler will be called when the upgrade
@@ -177,16 +180,15 @@ func (k Keeper) ScheduleUpgrade(ctx sdk.Context, plan types.Plan) error {
 		return types.ErrUpgradeCompleted
 	}
 
-	store := ctx.KVStore(k.storeKey)
-
 	// clear any old IBC state stored by previous plan
 	oldPlan, found := k.GetUpgradePlan(ctx)
 	if found {
-		k.ClearIBCState(ctx, oldPlan.Height)
+		for _, plan := range oldPlan {
+			k.ClearIBCState(ctx, plan.Height)
+		}
 	}
 
-	bz := k.cdc.MustMarshal(&plan)
-	store.Set(types.PlanKey(), bz)
+	k.upgradeConfig.SetPlan(&plan)
 
 	return nil
 }
@@ -282,13 +284,16 @@ func (k Keeper) ClearIBCState(ctx sdk.Context, lastHeight int64) {
 // ClearUpgradePlan clears any schedule upgrade and associated IBC states.
 func (k Keeper) ClearUpgradePlan(ctx sdk.Context) {
 	// clear IBC states everytime upgrade plan is removed
-	oldPlan, found := k.GetUpgradePlan(ctx)
+	planHeight := ctx.BlockHeight()
+	oldPlans, found := k.GetUpgradePlan(ctx)
 	if found {
-		k.ClearIBCState(ctx, oldPlan.Height)
+		for _, plan := range oldPlans {
+			planHeight = plan.Height
+			k.ClearIBCState(ctx, plan.Height)
+		}
 	}
 
-	store := ctx.KVStore(k.storeKey)
-	store.Delete(types.PlanKey())
+	k.upgradeConfig.ClearPlan(planHeight)
 }
 
 // Logger returns a module-specific logger.
@@ -298,17 +303,24 @@ func (k Keeper) Logger(ctx sdk.Context) log.Logger {
 
 // GetUpgradePlan returns the currently scheduled Plan if any, setting havePlan to true if there is a scheduled
 // upgrade or false if there is none
-func (k Keeper) GetUpgradePlan(ctx sdk.Context) (plan types.Plan, havePlan bool) {
-	plan, ok := k.upgradeConfig[ctx.BlockHeight()]
-	if !ok {
-		return types.Plan{}, false
+func (k Keeper) GetUpgradePlan(ctx sdk.Context) ([]*types.Plan, bool) {
+	plans := k.upgradeConfig.GetPlan(ctx.BlockHeight())
+	if len(plans) == 0 {
+		return nil, false
+	}
+
+	nonUpgraded := make([]*types.Plan, 0, len(plans))
+	for i := 0; i < len(plans); i++ {
+		if !k.IsUpgrade(ctx, plans[i].Name) {
+			nonUpgraded = append(nonUpgraded, plans[i])
+		}
 	}
 
-	if k.IsUpgrade(ctx, plan.Name) {
-		return types.Plan{}, false
+	if len(nonUpgraded) == 0 {
+		return nil, false
 	}
 
-	return plan, true
+	return nonUpgraded, true
 }
 
 // setDone marks this upgrade name as being done so the name can't be reused accidentally

x/upgrade/keeper/keeper_option.go

@@ -6,41 +6,51 @@ import (
 	"github.com/cosmos/cosmos-sdk/x/upgrade/types"
 )
 
-func convertUpgradeConfig(ctx sdk.Context, plans []serverconfig.UpgradeConfig) types.UpgradeConfig {
+func convertUpgradeConfig(ctx sdk.Context, plans []serverconfig.UpgradeConfig) (types.UpgradeConfig, error) {
 	upgradeConfig := types.NewUpgradeConfig()
 	if ctx.ChainID() == types.MainnetChainID {
 		upgradeConfig = types.MainnetConfig
 	}
 
 	// override by app config
 	for _, plan := range plans {
-		upgradeConfig.SetPlan(types.Plan{
+		nPlan := &types.Plan{
 			Name:   plan.Name,
 			Height: plan.Height,
 			Info:   plan.Info,
-		})
+		}
+		if err := nPlan.ValidateBasic(); err != nil {
+			return nil, err
+		}
+		upgradeConfig.SetPlan(nPlan)
 	}
 
-	return upgradeConfig
+	return upgradeConfig, nil
 }
 
 // Option function for Keeper
-type KeeperOption func(k *Keeper)
+type KeeperOption func(k *Keeper) error
 
 // RegisterUpgradePlan returns a KeeperOption to set the upgrade plan into the upgrade keeper
 func RegisterUpgradePlan(ctx sdk.Context,
 	plans []serverconfig.UpgradeConfig,
 ) KeeperOption {
-	return func(k *Keeper) {
-		k.upgradeConfig = convertUpgradeConfig(ctx, plans)
+	return func(k *Keeper) error {
+		c, err := convertUpgradeConfig(ctx, plans)
+		if err != nil {
+			return err
+		}
+		k.upgradeConfig = c
+		return nil
 	}
 }
 
 // RegisterUpgradeHandler returns a KeeperOption to set the upgrade handler into the upgrade keeper
 func RegisterUpgradeHandler(handlers map[string]types.UpgradeHandler) KeeperOption {
-	return func(k *Keeper) {
+	return func(k *Keeper) error {
 		for name, handler := range handlers {
 			k.SetUpgradeHandler(name, handler)
 		}
+		return nil
 	}
 }

x/upgrade/keeper/keeper_test.go

@@ -25,11 +25,15 @@ type KeeperTestSuite struct {
 }
 
 func (s *KeeperTestSuite) SetupTest() {
+	var err error
 	app := simapp.Setup(s.T(), false)
 	homeDir := filepath.Join(s.T().TempDir(), "x_upgrade_keeper_test")
-	app.UpgradeKeeper = keeper.NewKeeper( // recreate keeper in order to use a custom home path
+	app.UpgradeKeeper, err = keeper.NewKeeper( // recreate keeper in order to use a custom home Path
 		make(map[int64]bool), app.GetKey(types.StoreKey), app.AppCodec(), homeDir,
 	)
+	if err != nil {
+		s.T().Fatal(err)
+	}
 	s.T().Log("home dir:", homeDir)
 	s.homeDir = homeDir
 	s.app = app

x/upgrade/types/keys.go

@@ -17,8 +17,6 @@ const (
 )
 
 const (
-	// PlanByte specifies the Byte under which a pending upgrade plan is stored in the store
-	PlanByte = 0x0
 	// DoneByte is a prefix for to look up completed upgrade plan by name
 	DoneByte = 0x1
 
@@ -38,12 +36,6 @@ const (
 	KeyUpgradedConsState = "upgradedConsState"
 )
 
-// PlanKey is the key under which the current plan is saved
-// We store PlanByte as a const to keep it immutable (unlike a []byte)
-func PlanKey() []byte {
-	return []byte{PlanByte}
-}
-
 // UpgradedClientKey is the key under which the upgraded client state is saved
 // Connecting IBC chains can verify against the upgraded client in this path before
 // upgrading their clients