From 13125fab2840209484ca5fd1ec2e682651d1b109 Mon Sep 17 00:00:00 2001
From: Amaury Martiny <amaury.martiny@protonmail.com>
Date: Fri, 3 Jul 2020 19:18:13 +0200
Subject: [PATCH] Update tests

---
 x/upgrade/keeper/grpc_query.go      |  4 +--
 x/upgrade/keeper/grpc_query_test.go | 40 +++++++++++++++++++----------
 2 files changed, 28 insertions(+), 16 deletions(-)

diff --git a/x/upgrade/keeper/grpc_query.go b/x/upgrade/keeper/grpc_query.go
index 649019572f1a..ba7bd959c623 100644
--- a/x/upgrade/keeper/grpc_query.go
+++ b/x/upgrade/keeper/grpc_query.go
@@ -15,7 +15,7 @@ func (k Keeper) CurrentPlan(c context.Context, req *types.QueryCurrentPlanReques
 
 	plan, found := k.GetUpgradePlan(ctx)
 	if !found {
-		return nil, nil
+		return &types.QueryCurrentPlanResponse{}, nil
 	}
 
 	return &types.QueryCurrentPlanResponse{Plan: &plan}, nil
@@ -27,7 +27,7 @@ func (k Keeper) AppliedPlan(c context.Context, req *types.QueryAppliedPlanReques
 
 	applied := k.GetDoneHeight(ctx, req.Name)
 	if applied == 0 {
-		return nil, nil
+		return &types.QueryAppliedPlanResponse{}, nil
 	}
 
 	return &types.QueryAppliedPlanResponse{Height: applied}, nil
diff --git a/x/upgrade/keeper/grpc_query_test.go b/x/upgrade/keeper/grpc_query_test.go
index 2d3331314832..0e45d6e31d34 100644
--- a/x/upgrade/keeper/grpc_query_test.go
+++ b/x/upgrade/keeper/grpc_query_test.go
@@ -4,29 +4,41 @@ import (
 	gocontext "context"
 	"testing"
 
+	"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 TestGRPCQueryParams(t *testing.T) {
-	// TODO Setup tests - create queryClient
+func TestGRPCQueryUpgrade(t *testing.T) {
+	app := simapp.Setup(false)
+	ctx := app.BaseApp.NewContext(false, abci.Header{})
 
-	// No current plan
-	res, err := queryClient.Parameters(gocontext.Background(), &types.QueryCurrentPlanRequest{})
-	require.NoError(t, err)
-	require.Nil(t, res.Plan)
+	queryHelper := baseapp.NewQueryServerTestHelper(ctx)
+	types.RegisterQueryServer(queryHelper, app.UpgradeKeeper)
+	queryClient := types.NewQueryClient(queryHelper)
+
+	t.Log("Verify that the scheduled upgrade plan can be queried")
+	plan := types.Plan{Name: "test-plan", Height: 5}
+	app.UpgradeKeeper.ScheduleUpgrade(ctx, plan)
 
-	// TODO Create upgrade plan. Does it need to go through gov & upgrade proposal?
-	res, err := queryClient.Parameters(gocontext.Background(), &types.QueryCurrentPlanRequest{})
+	res, err := queryClient.CurrentPlan(gocontext.Background(), &types.QueryCurrentPlanRequest{})
 	require.NoError(t, err)
-	require.Equal(t, res.Plan, plan)
+	require.Equal(t, res.Plan, &plan)
 
-	// TODO Wait for plan to be applied. Possible to advance N blocks programmatically in test so that plan gets applied?
-	res, err := queryClient.Parameters(gocontext.Background(), &types.QueryCurrentPlanRequest{})
+	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 = queryClient.CurrentPlan(gocontext.Background(), &types.QueryCurrentPlanRequest{})
 	require.NoError(t, err)
 	require.Nil(t, res.Plan)
 
-	res, err := queryClient.Parameters(gocontext.Background(), &types.QueryAppliedPlanRequest{Name: "test-plan"})
-	require.NoError(t, err)
-	require.Equal(t, res.Height, 1234)
+	appliedRes, appliedErr := queryClient.AppliedPlan(gocontext.Background(), &types.QueryAppliedPlanRequest{Name: "test-plan"})
+
+	require.NoError(t, appliedErr)
+	require.Equal(t, int64(5), appliedRes.Height)
 }