From a3e6264416992c7e5e9cfb1ec904cca47caa7e94 Mon Sep 17 00:00:00 2001
From: mmsqe <mavis@crypto.com>
Date: Fri, 28 Mar 2025 22:13:47 +0800
Subject: [PATCH] fix(test): to avoid reusing previous store across upgrade
 test (#24170)

---
 x/upgrade/keeper/abci_test.go | 47 +++++++++++++++++------------------
 1 file changed, 23 insertions(+), 24 deletions(-)

diff --git a/x/upgrade/keeper/abci_test.go b/x/upgrade/keeper/abci_test.go
index 69589d264e..afe7fd698d 100644
--- a/x/upgrade/keeper/abci_test.go
+++ b/x/upgrade/keeper/abci_test.go
@@ -481,40 +481,20 @@ func TestBinaryVersion(t *testing.T) {
 }
 
 func TestDowngradeVerification(t *testing.T) {
-	s := setupTest(t, 10, map[int64]bool{})
-
-	// submit a plan.
 	planName := "downgrade"
-	height := s.env.HeaderService().HeaderInfo(s.ctx).Height
-	err := s.keeper.ScheduleUpgrade(s.ctx, types.Plan{Name: planName, Height: height + 1})
-	require.NoError(t, err)
-	s.ctx = s.ctx.WithHeaderInfo(header.Info{Height: height + 1})
-
-	// set the handler.
-	s.keeper.SetUpgradeHandler(planName, func(_ context.Context, _ types.Plan, vm appmodule.VersionMap) (appmodule.VersionMap, error) {
-		return vm, nil
-	})
-
-	// successful upgrade.
-	err = s.preModule.PreBlock(s.ctx)
-	require.NoError(t, err)
-
-	height = s.env.HeaderService().HeaderInfo(s.ctx).Height
-	s.ctx = s.ctx.WithHeaderInfo(header.Info{Height: height + 1})
-
 	testCases := map[string]struct {
-		preRun      func(*keeper.Keeper, context.Context, string)
+		preRun      func(*TestSuite, *keeper.Keeper, context.Context, string)
 		expectError bool
 	}{
 		"valid binary": {
-			preRun: func(k *keeper.Keeper, ctx context.Context, name string) {
+			preRun: func(_ *TestSuite, k *keeper.Keeper, ctx context.Context, name string) {
 				k.SetUpgradeHandler(planName, func(ctx context.Context, plan types.Plan, vm appmodule.VersionMap) (appmodule.VersionMap, error) {
 					return vm, nil
 				})
 			},
 		},
 		"downgrade with an active plan": {
-			preRun: func(k *keeper.Keeper, ctx context.Context, name string) {
+			preRun: func(s *TestSuite, k *keeper.Keeper, ctx context.Context, name string) {
 				height := s.env.HeaderService().HeaderInfo(s.ctx).Height
 				err := k.ScheduleUpgrade(ctx, types.Plan{Name: "another" + planName, Height: height + 1})
 				require.NoError(t, err, name)
@@ -527,6 +507,25 @@ func TestDowngradeVerification(t *testing.T) {
 	}
 
 	for name, tc := range testCases {
+		s := setupTest(t, 10, map[int64]bool{})
+		// submit a plan.
+		height := s.env.HeaderService().HeaderInfo(s.ctx).Height
+		err := s.keeper.ScheduleUpgrade(s.ctx, types.Plan{Name: planName, Height: height + 1})
+		require.NoError(t, err)
+		s.ctx = s.ctx.WithHeaderInfo(header.Info{Height: height + 1})
+
+		// set the handler.
+		s.keeper.SetUpgradeHandler(planName, func(_ context.Context, _ types.Plan, vm appmodule.VersionMap) (appmodule.VersionMap, error) {
+			return vm, nil
+		})
+
+		// successful upgrade.
+		err = s.preModule.PreBlock(s.ctx)
+		require.NoError(t, err)
+
+		height = s.env.HeaderService().HeaderInfo(s.ctx).Height
+		s.ctx = s.ctx.WithHeaderInfo(header.Info{Height: height + 1})
+
 		ctx := s.ctx
 
 		authority, err := addresscodec.NewBech32Codec("cosmos").BytesToString(authtypes.NewModuleAddress(govModuleName))
@@ -549,7 +548,7 @@ func TestDowngradeVerification(t *testing.T) {
 		require.ErrorIs(t, err, types.ErrNoUpgradePlanFound)
 
 		if tc.preRun != nil {
-			tc.preRun(k, ctx, name)
+			tc.preRun(s, k, ctx, name)
 		}
 
 		err = m.PreBlock(ctx)