x/upgrade: Refactor CLI to use protobuf query

x/upgrade/client/cli/query.go

@@ -1,98 +1,85 @@
 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"
 )
 
-// GetPlanCmd returns the query upgrade plan command
-func GetPlanCmd(storeName string, cdc *codec.Codec) *cobra.Command {
+// GetCurrentPlanCmd returns the query upgrade plan command
+func GetCurrentPlanCmd(clientCtx client.Context) *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)
+			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 the height at which a completed upgrade was applied
+func GetAppliedPlanCmd(clientCtx client.Context) *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)
+			queryClient := types.NewQueryClient(clientCtx.Init())
 
 			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 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/keeper/grpc_query_test.go

@@ -4,40 +4,61 @@ import (
 	gocontext "context"
 	"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 {
+	suite.Suite
+
+	app         *simapp.SimApp
+	ctx         sdk.Context
+	queryClient types.QueryClient
+}
+
+func (suite *UpgradeTestSuite) SetupTest() {
 	app := simapp.Setup(false)
 	ctx := app.BaseApp.NewContext(false, abci.Header{})
 
 	queryHelper := baseapp.NewQueryServerTestHelper(ctx, app.InterfaceRegistry())
 	types.RegisterQueryServer(queryHelper, app.UpgradeKeeper)
 	queryClient := types.NewQueryClient(queryHelper)
 
-	t.Log("Verify that the scheduled upgrade plan can be queried")
+	suite.app = app
+	suite.ctx = ctx
+	suite.queryClient = queryClient
+}
+
+func (suite *UpgradeTestSuite) TestGRPCQuery() {
+	queryClient := suite.queryClient
+
+	suite.T().Log("Verify that the scheduled upgrade plan can be queried")
 	plan := types.Plan{Name: "test-plan", Height: 5}
-	app.UpgradeKeeper.ScheduleUpgrade(ctx, plan)
+	suite.app.UpgradeKeeper.ScheduleUpgrade(suite.ctx, plan)
 
 	res, err := queryClient.CurrentPlan(gocontext.Background(), &types.QueryCurrentPlanRequest{})
-	require.NoError(t, err)
-	require.Equal(t, res.Plan, &plan)
+	suite.NoError(err)
+	suite.Equal(res.Plan, &plan)
 
-	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)
+	suite.T().Log("Verify that the upgrade plan can be successfully applied and queried")
+	suite.ctx = suite.ctx.WithBlockHeight(5)
+	suite.app.UpgradeKeeper.SetUpgradeHandler("test-plan", func(ctx sdk.Context, plan types.Plan) {})
+	suite.app.UpgradeKeeper.ApplyUpgrade(suite.ctx, plan)
 
 	res, err = queryClient.CurrentPlan(gocontext.Background(), &types.QueryCurrentPlanRequest{})
-	require.NoError(t, err)
-	require.Nil(t, res.Plan)
+	suite.NoError(err)
+	suite.Nil(res.Plan)
 
 	appliedRes, appliedErr := queryClient.AppliedPlan(gocontext.Background(), &types.QueryAppliedPlanRequest{Name: "test-plan"})
-	require.NoError(t, appliedErr)
-	require.Equal(t, int64(5), appliedRes.Height)
+	suite.NoError(appliedErr)
+	suite.Equal(int64(5), appliedRes.Height)
+}
+
+func TestUpgradeTestSuite(t *testing.T) {
+	suite.Run(t, new(UpgradeTestSuite))
 }

x/upgrade/module.go

@@ -60,8 +60,8 @@ func (AppModuleBasic) GetQueryCmd(clientCtx client.Context) *cobra.Command {
 		Short: "Querying commands for the upgrade module",
 	}
 	queryCmd.AddCommand(flags.GetCommands(
-		cli.GetPlanCmd(types.StoreKey, clientCtx.Codec),
-		cli.GetAppliedHeightCmd(types.StoreKey, clientCtx.Codec),
+		cli.GetCurrentPlanCmd(clientCtx),
+		cli.GetAppliedPlanCmd(clientCtx),
 	)...)
 
 	return queryCmd