feat(upgrade): migrate delegate (#120)

app/upgrades/v6/upgrade.go

@@ -7,12 +7,16 @@ import (
 	"github.com/cosmos/cosmos-sdk/types/module"
 	bankkeeper "github.com/cosmos/cosmos-sdk/x/bank/keeper"
 	banktypes "github.com/cosmos/cosmos-sdk/x/bank/types"
+	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
 	upgradetypes "github.com/cosmos/cosmos-sdk/x/upgrade/types"
+	"github.com/ethereum/go-ethereum/common"
 
 	"github.com/functionx/fx-core/v6/app/keepers"
 	crosschainkeeper "github.com/functionx/fx-core/v6/x/crosschain/keeper"
 	govtypes "github.com/functionx/fx-core/v6/x/gov/types"
 	layer2types "github.com/functionx/fx-core/v6/x/layer2/types"
+	migratekeeper "github.com/functionx/fx-core/v6/x/migrate/keeper"
+	fxstakingkeeper "github.com/functionx/fx-core/v6/x/staking/keeper"
 )
 
 func CreateUpgradeHandler(
@@ -109,3 +113,36 @@ func MigrateLayer2Module(ctx sdk.Context, layer2CrossChainKeeper crosschainkeepe
 		layer2CrossChainKeeper.AddBridgeToken(ctx, address, fxTokenDenom)
 	}
 }
+
+func AutoUndelegate(ctx sdk.Context, stakingKeeper fxstakingkeeper.Keeper) []stakingtypes.Delegation {
+	var delegations []stakingtypes.Delegation
+	stakingKeeper.IterateAllDelegations(ctx, func(delegation stakingtypes.Delegation) (stop bool) {
+		delegations = append(delegations, delegation)
+		delegator := sdk.MustAccAddressFromBech32(delegation.DelegatorAddress)
+		valAddress, err := sdk.ValAddressFromBech32(delegation.ValidatorAddress)
+		if err != nil {
+			panic(err)
+		}
+		if delegator.Equals(valAddress) {
+			return false
+		}
+		if _, err := stakingKeeper.Undelegate(ctx, delegator, valAddress, delegation.Shares); err != nil {
+			panic(err)
+		}
+		return false
+	})
+	return delegations
+}
+
+func ExportDelegate(ctx sdk.Context, migrateKeeper migratekeeper.Keeper, delegations []stakingtypes.Delegation) []stakingtypes.Delegation {
+	for i := 0; i < len(delegations); i++ {
+		delegation := delegations[i]
+		delegator := sdk.MustAccAddressFromBech32(delegation.DelegatorAddress)
+		if !migrateKeeper.HasMigratedDirectionTo(ctx, common.BytesToAddress(delegator.Bytes())) {
+			delegations = append(delegations[:i], delegations[i+1:]...)
+			i--
+			continue
+		}
+	}
+	return delegations
+}

app/upgrades/v6/upgrade_test.go

@@ -5,10 +5,11 @@ import (
 	"strings"
 	"testing"
 
+	sdkmath "cosmossdk.io/math"
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/stretchr/testify/suite"
-	abci "github.com/tendermint/tendermint/abci/types"
 	tmrand "github.com/tendermint/tendermint/libs/rand"
 	tmproto "github.com/tendermint/tendermint/proto/tendermint/types"
 
@@ -17,22 +18,24 @@ import (
 	"github.com/functionx/fx-core/v6/testutil/helpers"
 	fxtypes "github.com/functionx/fx-core/v6/types"
 	layer2types "github.com/functionx/fx-core/v6/x/layer2/types"
+	migratetypes "github.com/functionx/fx-core/v6/x/migrate/types"
 )
 
 type UpgradeTestSuite struct {
 	suite.Suite
 
-	app *app.App
-	ctx sdk.Context
+	app       *app.App
+	ctx       sdk.Context
+	valNumber int
 }
 
 func TestUpgradeTestSuite(t *testing.T) {
 	suite.Run(t, new(UpgradeTestSuite))
 }
 
 func (s *UpgradeTestSuite) SetupTest() {
-	valNumber := tmrand.Intn(5) + 5
-	valSet, valAccounts, valBalances := helpers.GenerateGenesisValidator(valNumber, sdk.Coins{})
+	s.valNumber = tmrand.Intn(5) + 5
+	valSet, valAccounts, valBalances := helpers.GenerateGenesisValidator(s.valNumber, sdk.Coins{})
 	s.app = helpers.SetupWithGenesisValSet(s.T(), valSet, valAccounts, valBalances...)
 	s.ctx = s.app.NewContext(false, tmproto.Header{
 		ChainID:         fxtypes.MainnetChainId,
@@ -46,22 +49,7 @@ func (s *UpgradeTestSuite) SetupTest() {
 }
 
 func (s *UpgradeTestSuite) CommitBlock(block int64) {
-	_ = s.app.Commit()
-	block--
-
-	ctx := s.ctx
-	nextBlockHeight := ctx.BlockHeight() + 1
-
-	_ = s.app.BeginBlock(abci.RequestBeginBlock{
-		Header: tmproto.Header{
-			Height: nextBlockHeight,
-		},
-	})
-	s.ctx = s.app.NewContext(false, ctx.BlockHeader())
-	s.ctx = s.ctx.WithBlockHeight(nextBlockHeight)
-	if block > 0 {
-		s.CommitBlock(block)
-	}
+	helpers.MintBlock(s.app, s.ctx, block)
 }
 
 func (s *UpgradeTestSuite) TestUpdateParams() {
@@ -96,3 +84,33 @@ func (s *UpgradeTestSuite) TestMigrateLayer2Module() {
 		s.Equal(bridgeToken.Denom, fmt.Sprintf("%s%s", layer2types.ModuleName, address))
 	}
 }
+
+func (s *UpgradeTestSuite) TestAutoUndelegate_And_ExportDelegate() {
+	delPrivKey := helpers.NewPriKey()
+	delAddr := sdk.AccAddress(delPrivKey.PubKey().Address())
+	helpers.AddTestAddr(s.app, s.ctx, delAddr, sdk.NewCoins(sdk.NewCoin(fxtypes.DefaultDenom, sdk.NewInt(10000))))
+	account := s.app.AccountKeeper.GetAccount(s.ctx, delAddr)
+	s.NoError(account.SetPubKey(delPrivKey.PubKey()))
+	s.app.AccountKeeper.SetAccount(s.ctx, account)
+
+	validators := s.app.StakingKeeper.GetAllValidators(s.ctx)
+	s.Equal(s.valNumber, len(validators))
+	validator := validators[0]
+
+	_, err := s.app.StakingKeeper.Delegate(s.ctx, delAddr, sdkmath.NewInt(100), stakingtypes.Unbonded, validator, true)
+	s.NoError(err)
+
+	newPrivKey := helpers.NewEthPrivKey()
+	_, err = s.app.MigrateKeeper.MigrateAccount(s.ctx, &migratetypes.MsgMigrateAccount{
+		From:      delAddr.String(),
+		To:        common.BytesToAddress(newPrivKey.PubKey().Address()).String(),
+		Signature: "",
+	})
+	s.NoError(err)
+
+	delegations := v6.AutoUndelegate(s.ctx, s.app.StakingKeeper)
+	s.Equal(s.valNumber+1, len(delegations))
+
+	delegations = v6.ExportDelegate(s.ctx, s.app.MigrateKeeper, delegations)
+	s.Equal(1, len(delegations))
+}