chore: add v4 upgrade handler

.changeset/entries/0578182a108ccf48ff35bf05cd2d09bdad445d6c40e41b6fe4f362af4270ce7e.yaml

@@ -0,0 +1,6 @@
+type: feat
+module: none
+pull_request: 222
+description: Added `v4` upgrade handler
+backward_compatible: true
+date: 2024-12-17T01:38:30.312655103Z

app/app.go

@@ -69,6 +69,7 @@ import (
 	"github.com/milkyway-labs/milkyway/v3/app/keepers"
 	"github.com/milkyway-labs/milkyway/v3/app/upgrades"
 	v3 "github.com/milkyway-labs/milkyway/v3/app/upgrades/v3"
+	v4 "github.com/milkyway-labs/milkyway/v3/app/upgrades/v4"
 	_ "github.com/milkyway-labs/milkyway/v3/client/docs/statik"
 	liquidvestingtypes "github.com/milkyway-labs/milkyway/v3/x/liquidvesting/types"
 )
@@ -79,6 +80,7 @@ var (
 
 	Upgrades = []upgrades.Upgrade{
 		v3.Upgrade,
+		v4.Upgrade,
 	}
 )
 

app/helpers/test_helpers.go

@@ -148,7 +148,8 @@ func setup() (*milkyway.MilkyWayApp, milkyway.GenesisState) {
 		appOptions,
 		emptyWasmOpts,
 	)
-	return milkyWayApp, milkyWayApp.ModuleBasics.DefaultGenesis(milkyWayApp.AppCodec())
+
+	return milkyWayApp, milkyway.NewDefaultGenesisState(milkyWayApp.AppCodec(), milkyWayApp.ModuleBasics)
 }
 
 func genesisStateWithValSet(t *testing.T,

app/upgrades/v4/constants.go

@@ -0,0 +1,18 @@
+package v4
+
+import (
+	storetypes "cosmossdk.io/store/types"
+
+	"github.com/milkyway-labs/milkyway/v3/app/upgrades"
+)
+
+const UpgradeName = "v4"
+
+var Upgrade = upgrades.Upgrade{
+	UpgradeName:          UpgradeName,
+	CreateUpgradeHandler: CreateUpgradeHandler,
+	StoreUpgrades: storetypes.StoreUpgrades{
+		Added:   []string{},
+		Deleted: []string{},
+	},
+}

app/upgrades/v4/upgrades.go

@@ -0,0 +1,52 @@
+package v4
+
+import (
+	"context"
+	"strings"
+
+	upgradetypes "cosmossdk.io/x/upgrade/types"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/module"
+
+	"github.com/milkyway-labs/milkyway/v3/app/keepers"
+)
+
+func CreateUpgradeHandler(
+	mm *module.Manager,
+	configuration module.Configurator,
+	keepers *keepers.AppKeepers,
+) upgradetypes.UpgradeHandler {
+	return func(ctx context.Context, plan upgradetypes.Plan, fromVM module.VersionMap) (module.VersionMap, error) {
+		vm, err := mm.RunMigrations(ctx, configuration, fromVM)
+		if err != nil {
+			return nil, err
+		}
+
+		// Unwrap the context
+		sdkCtx := sdk.UnwrapSDKContext(ctx)
+
+		// Remove all the markets except from the TIA/USD market
+		markets, err := keepers.MarketMapKeeper.GetAllMarkets(sdkCtx)
+		if err != nil {
+			return nil, err
+		}
+
+		for ticker, market := range markets {
+			if strings.Contains(ticker, "TIA") {
+				continue
+			}
+
+			err = keepers.MarketMapKeeper.DeleteMarket(sdkCtx, ticker)
+			if err != nil {
+				return nil, err
+			}
+
+			err = keepers.OracleKeeper.RemoveCurrencyPair(sdkCtx, market.Ticker.CurrencyPair)
+			if err != nil {
+				return nil, err
+			}
+		}
+
+		return vm, nil
+	}
+}

app/upgrades/v4/upgrades_test.go

@@ -0,0 +1,80 @@
+package v4_test
+
+import (
+	"testing"
+
+	"cosmossdk.io/core/appmodule"
+	"cosmossdk.io/core/header"
+	"cosmossdk.io/x/upgrade"
+	upgradetypes "cosmossdk.io/x/upgrade/types"
+	tmproto "github.com/cometbft/cometbft/proto/tendermint/types"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/stretchr/testify/suite"
+
+	milkywayapp "github.com/milkyway-labs/milkyway/v3/app"
+	"github.com/milkyway-labs/milkyway/v3/app/helpers"
+	v4 "github.com/milkyway-labs/milkyway/v3/app/upgrades/v4"
+)
+
+func TestUpgradeTestSuite(t *testing.T) {
+	suite.Run(t, new(UpgradeTestSuite))
+}
+
+type UpgradeTestSuite struct {
+	suite.Suite
+
+	App           *milkywayapp.MilkyWayApp
+	Ctx           sdk.Context
+	UpgradeModule appmodule.HasPreBlocker
+}
+
+func (suite *UpgradeTestSuite) SetupTest() {
+	suite.App = helpers.Setup(suite.T())
+	suite.Ctx = suite.App.NewUncachedContext(true, tmproto.Header{})
+	suite.Ctx = suite.Ctx.WithHeaderInfo(header.Info{Height: 1})
+	suite.UpgradeModule = upgrade.NewAppModule(suite.App.UpgradeKeeper, suite.App.AccountKeeper.AddressCodec())
+}
+
+func (suite *UpgradeTestSuite) TestUpgradeV4() {
+	// Make sure the markets are set
+	markets, err := suite.App.MarketMapKeeper.GetAllMarkets(suite.Ctx)
+	suite.Require().NoError(err)
+	suite.Require().Len(markets, 63)
+
+	// Make sure the currency pairs are set
+	currencyPairs := suite.App.OracleKeeper.GetAllCurrencyPairs(suite.Ctx)
+	suite.Require().Len(currencyPairs, 63)
+
+	// Perform the upgrade
+	suite.performUpgrade()
+
+	// Make sure only the TIA/USD market is left
+	markets, err = suite.App.MarketMapKeeper.GetAllMarkets(suite.Ctx)
+	suite.Require().NoError(err)
+	suite.Require().Len(markets, 1)
+	suite.Require().Equal("TIA/USD", markets["TIA/USD"].Ticker.String())
+
+	// Make sure the only currency pair left is TIA/USD
+	currencyPairs = suite.App.OracleKeeper.GetAllCurrencyPairs(suite.Ctx)
+	suite.Require().Len(currencyPairs, 1)
+	suite.Require().Equal("TIA/USD", currencyPairs[0].String())
+}
+
+func (suite *UpgradeTestSuite) performUpgrade() {
+	upgradeHeight := suite.Ctx.HeaderInfo().Height + 1
+
+	// Schedule the upgrade
+	err := suite.App.UpgradeKeeper.ScheduleUpgrade(suite.Ctx, upgradetypes.Plan{Name: v4.UpgradeName, Height: upgradeHeight})
+	suite.Require().NoError(err)
+
+	// Make sure the upgrade plan is set
+	_, err = suite.App.UpgradeKeeper.GetUpgradePlan(suite.Ctx)
+	suite.Require().NoError(err)
+
+	// Fast-forward to the upgrade height
+	suite.Ctx = suite.Ctx.WithHeaderInfo(header.Info{Height: upgradeHeight})
+
+	// Run PreBlocker to trigger the upgrade
+	_, err = suite.UpgradeModule.PreBlock(suite.Ctx)
+	suite.Require().NoError(err)
+}