switch from migration to upgrade

app/app.go

@@ -37,6 +37,7 @@ import (
 	servicemetrics "github.com/skip-mev/slinky/service/metrics"
 
 	v401 "github.com/neutron-org/neutron/v4/app/upgrades/v4.0.1"
+	v500 "github.com/neutron-org/neutron/v4/app/upgrades/v5.0.0"
 	"github.com/neutron-org/neutron/v4/x/globalfee"
 	globalfeetypes "github.com/neutron-org/neutron/v4/x/globalfee/types"
 
@@ -225,7 +226,7 @@ const (
 )
 
 var (
-	Upgrades = []upgrades.Upgrade{v401.Upgrade}
+	Upgrades = []upgrades.Upgrade{v401.Upgrade, v500.Upgrade}
 
 	// DefaultNodeHome default home directories for the application daemon
 	DefaultNodeHome string
@@ -1417,6 +1418,7 @@ func (app *App) setupUpgradeHandlers() {
 					MarketmapKeeper:     app.MarketMapKeeper,
 					FeeMarketKeeper:     app.FeeMarkerKeeper,
 					DynamicfeesKeeper:   app.DynamicFeesKeeper,
+					DexKeeper:           &app.DexKeeper,
 					GlobalFeeSubspace:   app.GetSubspace(globalfee.ModuleName),
 					CcvConsumerSubspace: app.GetSubspace(ccvconsumertypes.ModuleName),
 				},

app/upgrades/types.go

@@ -21,6 +21,7 @@ import (
 
 	contractmanagerkeeper "github.com/neutron-org/neutron/v4/x/contractmanager/keeper"
 	cronkeeper "github.com/neutron-org/neutron/v4/x/cron/keeper"
+	dexkeeper "github.com/neutron-org/neutron/v4/x/dex/keeper"
 	feeburnerkeeper "github.com/neutron-org/neutron/v4/x/feeburner/keeper"
 	icqkeeper "github.com/neutron-org/neutron/v4/x/interchainqueries/keeper"
 	tokenfactorykeeper "github.com/neutron-org/neutron/v4/x/tokenfactory/keeper"
@@ -64,6 +65,7 @@ type UpgradeKeepers struct {
 	MarketmapKeeper    *marketmapkeeper.Keeper
 	FeeMarketKeeper    *feemarketkeeper.Keeper
 	DynamicfeesKeeper  *dynamicfeeskeeper.Keeper
+	DexKeeper          *dexkeeper.Keeper
 	// subspaces
 	GlobalFeeSubspace   paramtypes.Subspace
 	CcvConsumerSubspace paramtypes.Subspace

app/upgrades/v5.0.0/constants.go

@@ -0,0 +1,20 @@
+package v400
+
+import (
+	storetypes "cosmossdk.io/store/types"
+
+	"github.com/neutron-org/neutron/v4/app/upgrades"
+)
+
+const (
+	// UpgradeName defines the on-chain upgrade name.
+	UpgradeName = "v5.0.0"
+)
+
+var Upgrade = upgrades.Upgrade{
+	UpgradeName:          UpgradeName,
+	CreateUpgradeHandler: CreateUpgradeHandler,
+	StoreUpgrades: storetypes.StoreUpgrades{
+		Added: []string{},
+	},
+}

x/dex/migrations/v5/store.go → app/upgrades/v5.0.0/upgrades.go

@@ -1,31 +1,48 @@
-package v5
+package v400
 
 import (
 	"context"
 	"fmt"
 
 	"cosmossdk.io/math"
+
+	upgradetypes "cosmossdk.io/x/upgrade/types"
+	"github.com/cosmos/cosmos-sdk/codec"
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/module"
 
-	"github.com/neutron-org/neutron/v4/x/dex/types"
+	"github.com/neutron-org/neutron/v4/app/upgrades"
+	dexkeeper "github.com/neutron-org/neutron/v4/x/dex/keeper"
 )
 
-type DexKeeperI interface {
-	GetAllPoolShareholders(sdk.Context) map[uint64][]types.PoolShareholder
-	GetPoolByID(sdk.Context, uint64) (*types.Pool, bool)
-	WithdrawCore(context.Context, *types.PairID, sdk.AccAddress, sdk.AccAddress, []math.Int, []int64, []uint64) error
-}
+func CreateUpgradeHandler(
+	mm *module.Manager,
+	configurator module.Configurator,
+	keepers *upgrades.UpgradeKeepers,
+	_ upgrades.StoreKeys,
+	_ codec.Codec,
+) upgradetypes.UpgradeHandler {
+	return func(c context.Context, _ upgradetypes.Plan, vm module.VersionMap) (module.VersionMap, error) {
+		ctx := sdk.UnwrapSDKContext(c)
 
-// MigrateStore performs in-place store migrations.
-func MigrateStore(ctx sdk.Context, k DexKeeperI) error {
-	if err := migratePools(ctx, k); err != nil {
-		return err
-	}
+		ctx.Logger().Info("Starting module migrations...")
+		vm, err := mm.RunMigrations(ctx, configurator, vm)
+		if err != nil {
+			return vm, err
+		}
 
-	return nil
+		ctx.Logger().Info("Running dex upgrades...")
+		err = upgradePools(ctx, *keepers.DexKeeper)
+		if err != nil {
+			return nil, err
+		}
+
+		ctx.Logger().Info(fmt.Sprintf("Migration {%s} applied", UpgradeName))
+		return vm, nil
+	}
 }
 
-func migratePools(ctx sdk.Context, k DexKeeperI) error {
+func upgradePools(ctx sdk.Context, k dexkeeper.Keeper) error {
 	// Due to an issue with autoswap logic any pools with multiple shareholders must be withdrawn to ensure correct accounting
 	ctx.Logger().Info("Migrating Pools...")
 

x/dex/migrations/v5/store_test.go → app/upgrades/v5.0.0/upgrades_test.go

@@ -1,36 +1,40 @@
-package v5_test
+package v400_test
 
 import (
 	"testing"
 
 	"cosmossdk.io/math"
-	"github.com/stretchr/testify/suite"
-
+	upgradetypes "cosmossdk.io/x/upgrade/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	bankkeeper "github.com/cosmos/cosmos-sdk/x/bank/keeper"
 
-	"github.com/neutron-org/neutron/v4/testutil"
+	v500 "github.com/neutron-org/neutron/v4/app/upgrades/v5.0.0"
 
-	bankkeeper "github.com/cosmos/cosmos-sdk/x/bank/keeper"
+	"github.com/stretchr/testify/suite"
 
+	"github.com/neutron-org/neutron/v4/testutil"
 	dexkeeper "github.com/neutron-org/neutron/v4/x/dex/keeper"
-	v5 "github.com/neutron-org/neutron/v4/x/dex/migrations/v5"
-	"github.com/neutron-org/neutron/v4/x/dex/types"
+	dextypes "github.com/neutron-org/neutron/v4/x/dex/types"
 )
 
-type V4DexMigrationTestSuite struct {
+type UpgradeTestSuite struct {
 	testutil.IBCConnectionTestSuite
 }
 
 func TestKeeperTestSuite(t *testing.T) {
-	suite.Run(t, new(V4DexMigrationTestSuite))
+	suite.Run(t, new(UpgradeTestSuite))
+}
+
+func (suite *UpgradeTestSuite) SetupTest() {
+	suite.IBCConnectionTestSuite.SetupTest()
 }
 
-func (suite *V4DexMigrationTestSuite) TestPoolMigrationSingleShareHolder() {
+func (suite *UpgradeTestSuite) TestPoolMigrationSingleShareHolder() {
 	var (
 		app           = suite.GetNeutronZoneApp(suite.ChainA)
 		ctx           = suite.ChainA.GetContext()
 		alice         = []byte("alice")
-		pairID        = &types.PairID{Token0: "TokenA", Token1: "TokenB"}
+		pairID        = &dextypes.PairID{Token0: "TokenA", Token1: "TokenB"}
 		depositAmount = math.NewInt(10_000)
 	)
 
@@ -39,11 +43,16 @@ func (suite *V4DexMigrationTestSuite) TestPoolMigrationSingleShareHolder() {
 	shares, err := suite.makeDeposit(ctx, app.DexKeeper, alice, pairID, depositAmount, math.ZeroInt(), 0, 1)
 	suite.NoError(err)
 
-	// run migrations
-	suite.NoError(v5.MigrateStore(ctx, app.DexKeeper))
+	// run upgrade
+	upgrade := upgradetypes.Plan{
+		Name:   v500.UpgradeName,
+		Info:   "some text here",
+		Height: 100,
+	}
+	suite.NoError(app.UpgradeKeeper.ApplyUpgrade(ctx, upgrade))
 
 	// assert pool and shareholder balance are unchanged
-	poolID, err := types.ParsePoolIDFromDenom(shares[0].Denom)
+	poolID, err := dextypes.ParsePoolIDFromDenom(shares[0].Denom)
 	suite.NoError(err)
 
 	pool, _ := app.DexKeeper.GetPoolByID(ctx, poolID)
@@ -53,13 +62,13 @@ func (suite *V4DexMigrationTestSuite) TestPoolMigrationSingleShareHolder() {
 	suite.True(aliceBalance.Equal(shares))
 }
 
-func (suite *V4DexMigrationTestSuite) TestPoolMigrationMultiShareHolder() {
+func (suite *UpgradeTestSuite) TestPoolMigrationMultiShareHolder() {
 	var (
 		app            = suite.GetNeutronZoneApp(suite.ChainA)
 		ctx            = suite.ChainA.GetContext()
 		alice          = []byte("alice")
 		bob            = []byte("bob")
-		pairID         = &types.PairID{Token0: "TokenA", Token1: "TokenB"}
+		pairID         = &dextypes.PairID{Token0: "TokenA", Token1: "TokenB"}
 		depositAmount  = math.NewInt(10_000)
 		initialBalance = sdk.NewCoins(sdk.NewCoin("TokenA", depositAmount))
 	)
@@ -77,11 +86,16 @@ func (suite *V4DexMigrationTestSuite) TestPoolMigrationMultiShareHolder() {
 	bobBalance := app.BankKeeper.GetAllBalances(ctx, bob)
 	suite.True(bobBalance.Equal(shares))
 
-	// run migrations
-	suite.NoError(v5.MigrateStore(ctx, app.DexKeeper))
+	// run upgrade
+	upgrade := upgradetypes.Plan{
+		Name:   v500.UpgradeName,
+		Info:   "some text here",
+		Height: 100,
+	}
+	suite.NoError(app.UpgradeKeeper.ApplyUpgrade(ctx, upgrade))
 
 	// assert that all users have withdrawn from the pool
-	poolID, err := types.ParsePoolIDFromDenom(shares[0].Denom)
+	poolID, err := dextypes.ParsePoolIDFromDenom(shares[0].Denom)
 	suite.NoError(err)
 
 	pool, _ := app.DexKeeper.GetPoolByID(ctx, poolID)
@@ -96,25 +110,25 @@ func (suite *V4DexMigrationTestSuite) TestPoolMigrationMultiShareHolder() {
 }
 
 func FundAccount(bankKeeper bankkeeper.Keeper, ctx sdk.Context, addr sdk.AccAddress, amounts sdk.Coins) {
-	if err := bankKeeper.MintCoins(ctx, types.ModuleName, amounts); err != nil {
+	if err := bankKeeper.MintCoins(ctx, dextypes.ModuleName, amounts); err != nil {
 		panic(err)
 	}
 
-	if err := bankKeeper.SendCoinsFromModuleToAccount(ctx, types.ModuleName, addr, amounts); err != nil {
+	if err := bankKeeper.SendCoinsFromModuleToAccount(ctx, dextypes.ModuleName, addr, amounts); err != nil {
 		panic(err)
 	}
 }
 
-func (suite *V4DexMigrationTestSuite) makeDeposit(
+func (suite *UpgradeTestSuite) makeDeposit(
 	ctx sdk.Context,
 	k dexkeeper.Keeper,
 	addr sdk.AccAddress,
-	pairID *types.PairID,
+	pairID *dextypes.PairID,
 	amount0, amount1 math.Int,
 	tick int64,
 	fee uint64,
 ) (sharesIssued sdk.Coins, err error) {
-	deposit0, deposit1, sharesIssued, _, err := k.DepositCore(ctx, pairID, addr, addr, []math.Int{amount0}, []math.Int{amount1}, []int64{tick}, []uint64{fee}, []*types.DepositOptions{{}})
+	deposit0, deposit1, sharesIssued, _, err := k.DepositCore(ctx, pairID, addr, addr, []math.Int{amount0}, []math.Int{amount1}, []int64{tick}, []uint64{fee}, []*dextypes.DepositOptions{{}})
 	suite.True(deposit0[0].Equal(amount0))
 	suite.True(deposit1[0].Equal(amount1))
 

x/dex/keeper/migrations.go

@@ -5,7 +5,6 @@ import (
 
 	v3 "github.com/neutron-org/neutron/v4/x/dex/migrations/v3"
 	v4 "github.com/neutron-org/neutron/v4/x/dex/migrations/v4"
-	v5 "github.com/neutron-org/neutron/v4/x/dex/migrations/v5"
 )
 
 // Migrator is a struct for handling in-place store migrations.
@@ -27,8 +26,3 @@ func (m Migrator) Migrate2to3(ctx sdk.Context) error {
 func (m Migrator) Migrate3to4(ctx sdk.Context) error {
 	return v4.MigrateStore(ctx, m.keeper.cdc, m.keeper.storeKey)
 }
-
-// Migrate4to5 migrates from version 4 to 5.
-func (m Migrator) Migrate4to5(ctx sdk.Context) error {
-	return v5.MigrateStore(ctx, m.keeper)
-}

x/dex/module.go

@@ -162,10 +162,6 @@ func (am AppModule) RegisterServices(cfg module.Configurator) {
 	if err := cfg.RegisterMigration(types.ModuleName, 3, m.Migrate3to4); err != nil {
 		panic(fmt.Sprintf("failed to migrate x/dex from version 3 to 4: %v", err))
 	}
-
-	if err := cfg.RegisterMigration(types.ModuleName, 4, m.Migrate4to5); err != nil {
-		panic(fmt.Sprintf("failed to migrate x/dex from version 4 to 5: %v", err))
-	}
 }
 
 // RegisterInvariants registers the capability module's invariants.

x/dex/types/constants.go

@@ -1,3 +1,3 @@
 package types
 
-const ConsensusVersion = 5
+const ConsensusVersion = 4