From 4bc8b488747afd6039ecd86917bdf43d27cd0c04 Mon Sep 17 00:00:00 2001
From: j75689 <j75689@gmail.com>
Date: Thu, 5 Jan 2023 17:51:18 +0800
Subject: [PATCH] fix: unit for upgrade module

---
 x/upgrade/abci_test.go              | 2 ++
 x/upgrade/keeper/grpc_query_test.go | 3 +++
 x/upgrade/keeper/keeper_test.go     | 6 ++++++
 3 files changed, 11 insertions(+)

diff --git a/x/upgrade/abci_test.go b/x/upgrade/abci_test.go
index fda44c6a87..caf50ec12b 100644
--- a/x/upgrade/abci_test.go
+++ b/x/upgrade/abci_test.go
@@ -124,6 +124,7 @@ func TestHaltIfTooNew(t *testing.T) {
 		called++
 		return vm, nil
 	})
+	s.keeper.SetUpgradeInitializer("future", func() error { return nil })
 
 	newCtx := s.ctx.WithBlockHeight(s.ctx.BlockHeight() + 1).WithBlockTime(time.Now())
 	req := abci.RequestBeginBlock{Header: newCtx.BlockHeader()}
@@ -201,6 +202,7 @@ func TestUpgrade(t *testing.T) {
 	s.keeper.SetUpgradeHandler("test", func(ctx sdk.Context, plan types.Plan, vm module.VersionMap) (module.VersionMap, error) {
 		return vm, nil
 	})
+	s.keeper.SetUpgradeInitializer("test", func() error { return nil })
 	require.NoError(t, err)
 	t.Log("Verify if upgrade happens without skip upgrade")
 	require.NotPanics(t, func() {
diff --git a/x/upgrade/keeper/grpc_query_test.go b/x/upgrade/keeper/grpc_query_test.go
index 9a0ab17f88..dd0c8eb489 100644
--- a/x/upgrade/keeper/grpc_query_test.go
+++ b/x/upgrade/keeper/grpc_query_test.go
@@ -117,6 +117,9 @@ func (suite *UpgradeTestSuite) TestAppliedCurrentPlan() {
 				suite.app.UpgradeKeeper.SetUpgradeHandler(planName, func(ctx sdk.Context, plan types.Plan, vm module.VersionMap) (module.VersionMap, error) {
 					return vm, nil
 				})
+				suite.app.UpgradeKeeper.SetUpgradeInitializer(planName, func() error {
+					return nil
+				})
 				suite.app.UpgradeKeeper.ApplyUpgrade(suite.ctx, plan)
 
 				req = &types.QueryAppliedPlanRequest{Name: planName}
diff --git a/x/upgrade/keeper/keeper_test.go b/x/upgrade/keeper/keeper_test.go
index 0993938d83..ee91bf1c31 100644
--- a/x/upgrade/keeper/keeper_test.go
+++ b/x/upgrade/keeper/keeper_test.go
@@ -125,6 +125,7 @@ func (s *KeeperTestSuite) TestScheduleUpgrade() {
 				s.app.UpgradeKeeper.SetUpgradeHandler("all-good", func(ctx sdk.Context, plan types.Plan, vm module.VersionMap) (module.VersionMap, error) {
 					return vm, nil
 				})
+				s.app.UpgradeKeeper.SetUpgradeInitializer("all-good", func() error { return nil })
 				s.app.UpgradeKeeper.ApplyUpgrade(s.ctx, types.Plan{
 					Name:   "all-good",
 					Info:   "some text here",
@@ -210,6 +211,7 @@ func (s *KeeperTestSuite) TestMigrations() {
 		vm["bank"] = vm["bank"] + 1
 		return vm, nil
 	})
+	s.app.UpgradeKeeper.SetUpgradeInitializer("dummy", func() error { return nil })
 	dummyPlan := types.Plan{
 		Name:   "dummy",
 		Info:   "some text here",
@@ -233,6 +235,7 @@ func (s *KeeperTestSuite) TestLastCompletedUpgrade() {
 	keeper.SetUpgradeHandler("test0", func(_ sdk.Context, _ types.Plan, vm module.VersionMap) (module.VersionMap, error) {
 		return vm, nil
 	})
+	keeper.SetUpgradeInitializer("test0", func() error { return nil })
 
 	keeper.ApplyUpgrade(s.ctx, types.Plan{
 		Name:   "test0",
@@ -247,6 +250,7 @@ func (s *KeeperTestSuite) TestLastCompletedUpgrade() {
 	keeper.SetUpgradeHandler("test1", func(_ sdk.Context, _ types.Plan, vm module.VersionMap) (module.VersionMap, error) {
 		return vm, nil
 	})
+	keeper.SetUpgradeInitializer("test1", func() error { return nil })
 
 	newCtx := s.ctx.WithBlockHeight(15)
 	keeper.ApplyUpgrade(newCtx, types.Plan{
@@ -270,6 +274,7 @@ func (s *KeeperTestSuite) TestLastCompletedUpgradeOrdering() {
 	keeper.SetUpgradeHandler("test-v0.9", func(_ sdk.Context, _ types.Plan, vm module.VersionMap) (module.VersionMap, error) {
 		return vm, nil
 	})
+	keeper.SetUpgradeInitializer("test-v0.9", func() error { return nil })
 
 	keeper.ApplyUpgrade(s.ctx, types.Plan{
 		Name:   "test-v0.9",
@@ -284,6 +289,7 @@ func (s *KeeperTestSuite) TestLastCompletedUpgradeOrdering() {
 	keeper.SetUpgradeHandler("test-v0.10", func(_ sdk.Context, _ types.Plan, vm module.VersionMap) (module.VersionMap, error) {
 		return vm, nil
 	})
+	keeper.SetUpgradeInitializer("test-v0.10", func() error { return nil })
 
 	newCtx := s.ctx.WithBlockHeight(15)
 	keeper.ApplyUpgrade(newCtx, types.Plan{