x/upgrade: Refactor CLI to use protobuf query

x/upgrade/client/cli/query.go

@@ -1,98 +1,102 @@
 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
-func GetPlanCmd(storeName string, cdc *codec.Codec) *cobra.Command {
+// GetQueryCmd returns the parent command for all x/upgrade CLi query commands
+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.Init())
 
-			// ignore height for now
-			res, _, err := clientCtx.Query(fmt.Sprintf("custom/%s/%s", types.QuerierKey, types.QueryCurrent))
+			params := types.NewQueryCurrentPlanRequest()
+			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
-func GetAppliedHeightCmd(storeName string, cdc *codec.Codec) *cobra.Command {
+// GetAppliedPlanCmd returns information about the block at which a completed
+// 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",
 		Long: "If upgrade-name was previously executed on the chain, this returns the header for the block at which it was applied.\n" +
 			"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.Init())
 
 			name := args[0]
 			params := types.NewQueryAppliedPlanRequest(name)
-			bz, err := clientCtx.JSONMarshaler.MarshalJSON(params)
-			if err != nil {
-				return err
-			}
-
-			res, _, err := clientCtx.QueryWithData(fmt.Sprintf("custom/%s/%s", types.QuerierKey, types.QueryApplied), bz)
+			res, err := queryClient.AppliedPlan(context.Background(), params)
 			if err != nil {
 				return err
 			}
 
-			if len(res) == 0 {
+			if res.Height == 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 nil
+			return clientCtx.PrintOutput(string(bz))
 		},
 	}
 }

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) {
-	app := simapp.Setup(false)
-	ctx := app.BaseApp.NewContext(false, abci.Header{})
+type UpgradeTestSuite struct {
+	suite.Suite
+
+	app         *simapp.SimApp
+	ctx         sdk.Context
+	queryClient types.QueryClient
+}
+
+func (suite *UpgradeTestSuite) SetupTest() {
+	suite.app = simapp.Setup(false)
+	suite.ctx = suite.app.BaseApp.NewContext(false, abci.Header{})
+
+	queryHelper := baseapp.NewQueryServerTestHelper(suite.ctx, suite.app.InterfaceRegistry())
+	types.RegisterQueryServer(queryHelper, suite.app.UpgradeKeeper)
+	suite.queryClient = types.NewQueryClient(queryHelper)
+}
+
+func (suite *UpgradeTestSuite) TestQueryCurrentPlan() {
+	var (
+		req         *types.QueryCurrentPlanRequest
+		expResponse types.QueryCurrentPlanResponse
+	)
+
+	testCases := []struct {
+		msg      string
+		malleate func()
+		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)
 
-	queryHelper := baseapp.NewQueryServerTestHelper(ctx, app.InterfaceRegistry())
-	types.RegisterQueryServer(queryHelper, app.UpgradeKeeper)
-	queryClient := types.NewQueryClient(queryHelper)
+				req = types.NewQueryCurrentPlanRequest()
+				expResponse = types.QueryCurrentPlanResponse{Plan: &plan}
+			},
+			true,
+		},
+	}
 
-	t.Log("Verify that the scheduled upgrade plan can be queried")
-	plan := types.Plan{Name: "test-plan", Height: 5}
-	app.UpgradeKeeper.ScheduleUpgrade(ctx, plan)
+	for _, tc := range testCases {
+		suite.Run(fmt.Sprintf("Case %s", tc.msg), func() {
+			suite.SetupTest() // reset
 
-	res, err := queryClient.CurrentPlan(gocontext.Background(), &types.QueryCurrentPlanRequest{})
-	require.NoError(t, err)
-	require.Equal(t, res.Plan, &plan)
+			tc.malleate()
 
-	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) {})
-	app.UpgradeKeeper.ApplyUpgrade(ctx, plan)
+			res, err := suite.queryClient.CurrentPlan(gocontext.Background(), req)
 
-	res, err = queryClient.CurrentPlan(gocontext.Background(), &types.QueryCurrentPlanRequest{})
-	require.NoError(t, err)
-	require.Nil(t, res.Plan)
+			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)
+			}
+		})
+	}
+}
 
-	appliedRes, appliedErr := queryClient.AppliedPlan(gocontext.Background(), &types.QueryAppliedPlanRequest{Name: "test-plan"})
-	require.NoError(t, appliedErr)
-	require.Equal(t, int64(5), appliedRes.Height)
+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 {
-	queryCmd := &cobra.Command{
-		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
+func (AppModuleBasic) GetQueryCmd(_ client.Context) *cobra.Command {
+	return cli.GetQueryCmd()
 }
 
 // GetTxCmd returns the transaction commands for this module
 func (AppModuleBasic) GetTxCmd(_ client.Context) *cobra.Command {
-	txCmd := &cobra.Command{
-		Use:   "upgrade",
-		Short: "Upgrade transaction subcommands",
-	}
-	txCmd.AddCommand(flags.PostCommands()...)
-	return txCmd
+	return cli.GetTxCmd()
 }
 
 func (b AppModuleBasic) RegisterInterfaceTypes(registry codectypes.InterfaceRegistry) {