x/upgrade: Refactor CLI to use protobuf query (#6623)

* x/upgrade: Refactor CLI to use protobuf query

* Import lint

* Use table tests

* Small tweak in setup

* Address bez cli refactor

* Address fede's review

* Remove useless func args

* Add back clientCtx to tx command

* Update comments

* Update docs

* Small refactor

* remove Init()

Co-authored-by: Federico Kunze <31522760+fedekunze@users.noreply.github.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

x/upgrade/client/cli/query.go

@@ -1,49 +1,62 @@
 package cli
 
 import (
-	"encoding/binary"
+	"context"
 	"fmt"
 
-	"github.com/cosmos/cosmos-sdk/x/upgrade/types"
-
 	"github.com/spf13/cobra"
 
 	"github.com/cosmos/cosmos-sdk/client"
-	"github.com/cosmos/cosmos-sdk/codec"
+	"github.com/cosmos/cosmos-sdk/client/flags"
+	"github.com/cosmos/cosmos-sdk/x/upgrade/types"
 )
 
-// GetPlanCmd returns the query upgrade plan command
+// GetQueryCmd returns the parent command for all x/upgrade CLi query commands
-func GetPlanCmd(storeName string, cdc *codec.Codec) *cobra.Command {
+func GetQueryCmd() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   types.ModuleName,
+		Short: "Querying commands for the upgrade module",
+	}
+	cmd.AddCommand(flags.GetCommands(
+		GetCurrentPlanCmd(),
+		GetAppliedPlanCmd(),
+	)...)
+
+	return cmd
+}
+
+// GetCurrentPlanCmd returns the query upgrade plan command
+func GetCurrentPlanCmd() *cobra.Command {
 	return &cobra.Command{
 		Use:   "plan",
 		Short: "get upgrade plan (if one exists)",
 		Long:  "Gets the currently scheduled upgrade plan, if one exists",
 		Args:  cobra.ExactArgs(0),
 		RunE: func(cmd *cobra.Command, args []string) error {
-			clientCtx := client.NewContext().WithCodec(cdc).WithJSONMarshaler(cdc)
+			clientCtx := client.GetClientContextFromCmd(cmd)
+			queryClient := types.NewQueryClient(clientCtx)
 
-			// ignore height for now
+			params := types.NewQueryCurrentPlanRequest()
-			res, _, err := clientCtx.Query(fmt.Sprintf("custom/%s/%s", types.QuerierKey, types.QueryCurrent))
+			res, err := queryClient.CurrentPlan(context.Background(), params)
 			if err != nil {
 				return err
 			}
 
-			if len(res) == 0 {
+			if len(res.Plan.Name) == 0 {
 				return fmt.Errorf("no upgrade scheduled")
 			}
 
-			var plan types.Plan
-			err = cdc.UnmarshalJSON(res, &plan)
 			if err != nil {
 				return err
 			}
-			return clientCtx.PrintOutput(plan)
+			return clientCtx.PrintOutput(res.Plan)
 		},
 	}
 }
 
-// GetAppliedHeightCmd returns the height at which a completed upgrade was applied
+// GetAppliedPlanCmd returns information about the block at which a completed
-func GetAppliedHeightCmd(storeName string, cdc *codec.Codec) *cobra.Command {
+// upgrade was applied
+func GetAppliedPlanCmd() *cobra.Command {
 	return &cobra.Command{
 		Use:   "applied [upgrade-name]",
 		Short: "block header for height at which a completed upgrade was applied",
@@ -51,48 +64,39 @@ func GetAppliedHeightCmd(storeName string, cdc *codec.Codec) *cobra.Command {
 			"This helps a client determine which binary was valid over a given range of blocks, as well as more context to understand past migrations.",
 		Args: cobra.ExactArgs(1),
 		RunE: func(cmd *cobra.Command, args []string) error {
-			clientCtx := client.NewContext().WithCodec(cdc).WithJSONMarshaler(cdc)
+			clientCtx := client.GetClientContextFromCmd(cmd)
+			queryClient := types.NewQueryClient(clientCtx)
 
 			name := args[0]
 			params := types.NewQueryAppliedPlanRequest(name)
-			bz, err := clientCtx.JSONMarshaler.MarshalJSON(params)
+			res, err := queryClient.AppliedPlan(context.Background(), params)
 			if err != nil {
 				return err
 			}
 
-			res, _, err := clientCtx.QueryWithData(fmt.Sprintf("custom/%s/%s", types.QuerierKey, types.QueryApplied), bz)
+			if res.Height == 0 {
-			if err != nil {
-				return err
-			}
-
-			if len(res) == 0 {
 				return fmt.Errorf("no upgrade found")
 			}
-			if len(res) != 8 {
-				return fmt.Errorf("unknown format for applied-upgrade")
-			}
-			applied := int64(binary.BigEndian.Uint64(res))
 
 			// we got the height, now let's return the headers
 			node, err := clientCtx.GetNode()
 			if err != nil {
 				return err
 			}
-			headers, err := node.BlockchainInfo(applied, applied)
+			headers, err := node.BlockchainInfo(res.Height, res.Height)
 			if err != nil {
 				return err
 			}
 			if len(headers.BlockMetas) == 0 {
-				return fmt.Errorf("no headers returned for height %d", applied)
+				return fmt.Errorf("no headers returned for height %d", res.Height)
 			}
 
 			// always output json as Header is unreable in toml ([]byte is a long list of numbers)
-			bz, err = cdc.MarshalJSONIndent(headers.BlockMetas[0], "", "  ")
+			bz, err := clientCtx.Codec.MarshalJSONIndent(headers.BlockMetas[0], "", "  ")
 			if err != nil {
 				return err
 			}
-			fmt.Println(string(bz))
+			return clientCtx.PrintOutput(string(bz))
-			return nil
 		},
 	}
 }

x/upgrade/client/cli/tx.go

@@ -4,15 +4,14 @@ import (
 	"fmt"
 	"time"
 
-	"github.com/cosmos/cosmos-sdk/client/tx"
-	gov "github.com/cosmos/cosmos-sdk/x/gov/types"
-
-	"github.com/cosmos/cosmos-sdk/x/gov/client/cli"
-
 	"github.com/spf13/cobra"
 
 	"github.com/cosmos/cosmos-sdk/client"
+	"github.com/cosmos/cosmos-sdk/client/flags"
+	"github.com/cosmos/cosmos-sdk/client/tx"
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/x/gov/client/cli"
+	gov "github.com/cosmos/cosmos-sdk/x/gov/types"
 	"github.com/cosmos/cosmos-sdk/x/upgrade/types"
 )
 
@@ -25,9 +24,19 @@ const (
 	FlagUpgradeInfo   = "info"
 )
 
+// GetTxCmd returns the transaction commands for this module
+func GetTxCmd() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   types.ModuleName,
+		Short: "Upgrade transaction subcommands",
+	}
+	cmd.AddCommand(flags.PostCommands()...)
+
+	return cmd
+}
+
 // NewCmdSubmitUpgradeProposal implements a command handler for submitting a software upgrade proposal transaction.
 func NewCmdSubmitUpgradeProposal(clientCtx client.Context) *cobra.Command {
-
 	cmd := &cobra.Command{
 		Use:   "software-upgrade [name] (--upgrade-height [height] | --upgrade-time [time]) (--upgrade-info [info]) [flags]",
 		Args:  cobra.ExactArgs(1),

x/upgrade/keeper/grpc_query_test.go

@@ -2,42 +2,143 @@ package keeper_test
 
 import (
 	gocontext "context"
+	"fmt"
 	"testing"
 
+	"github.com/stretchr/testify/suite"
+
+	abci "github.com/tendermint/tendermint/abci/types"
+
 	"github.com/cosmos/cosmos-sdk/baseapp"
 	"github.com/cosmos/cosmos-sdk/simapp"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/x/upgrade/types"
-	"github.com/stretchr/testify/require"
-	abci "github.com/tendermint/tendermint/abci/types"
 )
 
-func TestGRPCQueryUpgrade(t *testing.T) {
+type UpgradeTestSuite struct {
-	app := simapp.Setup(false)
+	suite.Suite
-	ctx := app.BaseApp.NewContext(false, abci.Header{})
 
-	queryHelper := baseapp.NewQueryServerTestHelper(ctx, app.InterfaceRegistry())
+	app         *simapp.SimApp
-	types.RegisterQueryServer(queryHelper, app.UpgradeKeeper)
+	ctx         sdk.Context
-	queryClient := types.NewQueryClient(queryHelper)
+	queryClient types.QueryClient
-
+}
-	t.Log("Verify that the scheduled upgrade plan can be queried")
+
-	plan := types.Plan{Name: "test-plan", Height: 5}
+func (suite *UpgradeTestSuite) SetupTest() {
-	app.UpgradeKeeper.ScheduleUpgrade(ctx, plan)
+	suite.app = simapp.Setup(false)
-
+	suite.ctx = suite.app.BaseApp.NewContext(false, abci.Header{})
-	res, err := queryClient.CurrentPlan(gocontext.Background(), &types.QueryCurrentPlanRequest{})
+
-	require.NoError(t, err)
+	queryHelper := baseapp.NewQueryServerTestHelper(suite.ctx, suite.app.InterfaceRegistry())
-	require.Equal(t, res.Plan, &plan)
+	types.RegisterQueryServer(queryHelper, suite.app.UpgradeKeeper)
-
+	suite.queryClient = types.NewQueryClient(queryHelper)
-	t.Log("Verify that the upgrade plan can be successfully applied and queried")
+}
-	ctx = ctx.WithBlockHeight(5)
+
-	app.UpgradeKeeper.SetUpgradeHandler("test-plan", func(ctx sdk.Context, plan types.Plan) {})
+func (suite *UpgradeTestSuite) TestQueryCurrentPlan() {
-	app.UpgradeKeeper.ApplyUpgrade(ctx, plan)
+	var (
-
+		req         *types.QueryCurrentPlanRequest
-	res, err = queryClient.CurrentPlan(gocontext.Background(), &types.QueryCurrentPlanRequest{})
+		expResponse types.QueryCurrentPlanResponse
-	require.NoError(t, err)
+	)
-	require.Nil(t, res.Plan)
+
-
+	testCases := []struct {
-	appliedRes, appliedErr := queryClient.AppliedPlan(gocontext.Background(), &types.QueryAppliedPlanRequest{Name: "test-plan"})
+		msg      string
-	require.NoError(t, appliedErr)
+		malleate func()
-	require.Equal(t, int64(5), appliedRes.Height)
+		expPass  bool
+	}{
+		{
+			"without current upgrade plan",
+			func() {
+				req = types.NewQueryCurrentPlanRequest()
+				expResponse = types.QueryCurrentPlanResponse{}
+			},
+			true,
+		},
+		{
+			"with current upgrade plan",
+			func() {
+				plan := types.Plan{Name: "test-plan", Height: 5}
+				suite.app.UpgradeKeeper.ScheduleUpgrade(suite.ctx, plan)
+
+				req = types.NewQueryCurrentPlanRequest()
+				expResponse = types.QueryCurrentPlanResponse{Plan: &plan}
+			},
+			true,
+		},
+	}
+
+	for _, tc := range testCases {
+		suite.Run(fmt.Sprintf("Case %s", tc.msg), func() {
+			suite.SetupTest() // reset
+
+			tc.malleate()
+
+			res, err := suite.queryClient.CurrentPlan(gocontext.Background(), req)
+
+			if tc.expPass {
+				suite.Require().NoError(err)
+				suite.Require().NotNil(res)
+				suite.Require().Equal(&expResponse, res)
+			} else {
+				suite.Require().Error(err)
+			}
+		})
+	}
+}
+
+func (suite *UpgradeTestSuite) TestAppliedCurrentPlan() {
+	var (
+		req       *types.QueryAppliedPlanRequest
+		expHeight int64
+	)
+
+	testCases := []struct {
+		msg      string
+		malleate func()
+		expPass  bool
+	}{
+		{
+			"with non-existent upgrade plan",
+			func() {
+				req = types.NewQueryAppliedPlanRequest("foo")
+			},
+			true,
+		},
+		{
+			"with applied upgrade plan",
+			func() {
+				expHeight = 5
+
+				planName := "test-plan"
+				plan := types.Plan{Name: planName, Height: expHeight}
+				suite.app.UpgradeKeeper.ScheduleUpgrade(suite.ctx, plan)
+
+				suite.ctx = suite.ctx.WithBlockHeight(expHeight)
+				suite.app.UpgradeKeeper.SetUpgradeHandler(planName, func(ctx sdk.Context, plan types.Plan) {})
+				suite.app.UpgradeKeeper.ApplyUpgrade(suite.ctx, plan)
+
+				req = types.NewQueryAppliedPlanRequest(planName)
+			},
+			true,
+		},
+	}
+
+	for _, tc := range testCases {
+		suite.Run(fmt.Sprintf("Case %s", tc.msg), func() {
+			suite.SetupTest() // reset
+
+			tc.malleate()
+
+			res, err := suite.queryClient.AppliedPlan(gocontext.Background(), req)
+
+			if tc.expPass {
+				suite.Require().NoError(err)
+				suite.Require().NotNil(res)
+				suite.Require().Equal(expHeight, res.Height)
+			} else {
+				suite.Require().Error(err)
+			}
+		})
+	}
+}
+
+func TestUpgradeTestSuite(t *testing.T) {
+	suite.Run(t, new(UpgradeTestSuite))
 }

x/upgrade/module.go

@@ -11,7 +11,6 @@ import (
 	abci "github.com/tendermint/tendermint/abci/types"
 
 	"github.com/cosmos/cosmos-sdk/client"
-	"github.com/cosmos/cosmos-sdk/client/flags"
 	"github.com/cosmos/cosmos-sdk/codec"
 	codectypes "github.com/cosmos/cosmos-sdk/codec/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
@@ -54,27 +53,13 @@ func (AppModuleBasic) RegisterRESTRoutes(clientCtx client.Context, r *mux.Router
 }
 
 // GetQueryCmd returns the cli query commands for this module
-func (AppModuleBasic) GetQueryCmd(clientCtx client.Context) *cobra.Command {
+func (AppModuleBasic) GetQueryCmd(_ client.Context) *cobra.Command {
-	queryCmd := &cobra.Command{
+	return cli.GetQueryCmd()
-		Use:   "upgrade",
-		Short: "Querying commands for the upgrade module",
-	}
-	queryCmd.AddCommand(flags.GetCommands(
-		cli.GetPlanCmd(types.StoreKey, clientCtx.Codec),
-		cli.GetAppliedHeightCmd(types.StoreKey, clientCtx.Codec),
-	)...)
-
-	return queryCmd
 }
 
 // GetTxCmd returns the transaction commands for this module
 func (AppModuleBasic) GetTxCmd(_ client.Context) *cobra.Command {
-	txCmd := &cobra.Command{
+	return cli.GetTxCmd()
-		Use:   "upgrade",
-		Short: "Upgrade transaction subcommands",
-	}
-	txCmd.AddCommand(flags.PostCommands()...)
-	return txCmd
 }
 
 func (b AppModuleBasic) RegisterInterfaceTypes(registry codectypes.InterfaceRegistry) {