fix: verify delegation and reward data after upgrade (#943)

app/upgrade_test.go

@@ -2,6 +2,7 @@ package app_test
 
 import (
 	"errors"
+	"fmt"
 	"math/big"
 	"os"
 	"path/filepath"
@@ -21,11 +22,14 @@ import (
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
 	banktypes "github.com/cosmos/cosmos-sdk/x/bank/types"
+	distributionkeeper "github.com/cosmos/cosmos-sdk/x/distribution/keeper"
 	distributiontypes "github.com/cosmos/cosmos-sdk/x/distribution/types"
 	govtypes "github.com/cosmos/cosmos-sdk/x/gov/types"
+	stakingkeeper "github.com/cosmos/cosmos-sdk/x/staking/keeper"
 	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
 	"github.com/ethereum/go-ethereum/common"
 	evmtypes "github.com/evmos/ethermint/x/evm/types"
+	"github.com/rs/zerolog"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
@@ -51,7 +55,7 @@ func Test_UpgradeAndMigrate(t *testing.T) {
 
 	ctx := newContext(t, myApp, chainId, false)
 
-	bdd := beforeUpgrade(ctx, myApp)
+	bdd := beforeUpgrade(t, ctx, myApp)
 
 	// 1. set upgrade plan
 	require.NoError(t, myApp.UpgradeKeeper.ScheduleUpgrade(ctx, upgradetypes.Plan{
@@ -111,7 +115,7 @@ func buildApp(t *testing.T) *app.App {
 	require.NoError(t, err)
 
 	myApp := helpers.NewApp(func(opts *helpers.AppOpts) {
-		opts.Logger = log.NewLogger(os.Stdout)
+		opts.Logger = log.NewLogger(os.Stdout, log.LevelOption(zerolog.InfoLevel))
 		opts.DB = db
 		opts.Home = home
 	})
@@ -150,24 +154,63 @@ func newContext(t *testing.T, myApp *app.App, chainId string, deliveState bool)
 type BeforeUpgradeData struct {
 	AccountBalances map[string]sdk.Coins
 	ModuleBalances  map[string]sdk.Coins
+	Delegates       map[string]*stakingtypes.DelegationResponse
+	DelegateRewards map[string]sdk.DecCoins
 }
 
-func beforeUpgrade(ctx sdk.Context, myApp *app.App) BeforeUpgradeData {
+func beforeUpgrade(t *testing.T, ctx sdk.Context, myApp *app.App) BeforeUpgradeData {
+	t.Helper()
 	bdd := BeforeUpgradeData{
 		AccountBalances: make(map[string]sdk.Coins),
 		ModuleBalances:  make(map[string]sdk.Coins),
+		Delegates:       make(map[string]*stakingtypes.DelegationResponse),
+		DelegateRewards: make(map[string]sdk.DecCoins),
 	}
 
 	accountBalance, moduleBalance := allBalances(ctx, myApp)
 	bdd.AccountBalances = accountBalance
 	bdd.ModuleBalances = moduleBalance
 
+	delegatesMap, delegateRewardsMap := delegatesAndRewards(t, ctx, myApp)
+	bdd.Delegates = delegatesMap
+	bdd.DelegateRewards = delegateRewardsMap
 	return bdd
 }
 
+func delegatesAndRewards(t *testing.T, ctx sdk.Context, myApp *app.App) (map[string]*stakingtypes.DelegationResponse, map[string]sdk.DecCoins) {
+	t.Helper()
+
+	delegatesMap := make(map[string]*stakingtypes.DelegationResponse)
+	delegateRewardsMap := make(map[string]sdk.DecCoins)
+	distrQuerier := distributionkeeper.NewQuerier(myApp.DistrKeeper)
+	stakingQuerier := stakingkeeper.NewQuerier(myApp.StakingKeeper.Keeper)
+
+	err := myApp.StakingKeeper.IterateAllDelegations(ctx, func(del stakingtypes.Delegation) (stop bool) {
+		delegateKey := fmt.Sprintf("%s:%s", del.DelegatorAddress, del.ValidatorAddress)
+		delegation, err := stakingQuerier.Delegation(ctx, &stakingtypes.QueryDelegationRequest{
+			DelegatorAddr: del.DelegatorAddress,
+			ValidatorAddr: del.ValidatorAddress,
+		})
+		require.NoError(t, err)
+		delegatesMap[delegateKey] = delegation.DelegationResponse
+		cacheCtx, _ := ctx.CacheContext()
+		rewards, err := distrQuerier.DelegationRewards(cacheCtx, &distributiontypes.QueryDelegationRewardsRequest{
+			DelegatorAddress: del.DelegatorAddress,
+			ValidatorAddress: del.ValidatorAddress,
+		})
+		require.NoError(t, err)
+		delegateRewardsMap[delegateKey] = rewards.Rewards
+		return false
+	})
+	require.NoError(t, err)
+	return delegatesMap, delegateRewardsMap
+}
+
 func checkAppUpgrade(t *testing.T, ctx sdk.Context, myApp *app.App, bdd BeforeUpgradeData) {
 	t.Helper()
 
+	checkDelegationData(t, bdd, ctx, myApp)
+
 	checkWrapToken(t, ctx, myApp)
 
 	checkStakingMigrationDelete(t, ctx, myApp)
@@ -191,6 +234,26 @@ func checkAppUpgrade(t *testing.T, ctx sdk.Context, myApp *app.App, bdd BeforeUp
 	checkModulesData(t, ctx, myApp)
 }
 
+func checkDelegationData(t *testing.T, bdd BeforeUpgradeData, ctx sdk.Context, myApp *app.App) {
+	t.Helper()
+
+	beforeDelegates, beforeDelegateRewards := bdd.Delegates, bdd.DelegateRewards
+	afterDelegates, afterDelegateRewards := delegatesAndRewards(t, ctx, myApp)
+
+	for delegateKey, beforeDelegate := range beforeDelegates {
+		afterDelegate, ok := afterDelegates[delegateKey]
+		require.True(t, ok)
+		assert.EqualValues(t, beforeDelegate.Delegation.String(), afterDelegate.Delegation.String(), delegateKey)
+		assert.EqualValues(t, fxtypes.SwapCoin(beforeDelegate.Balance).String(), afterDelegate.Balance.String(), delegateKey, beforeDelegate.Balance.String())
+	}
+	for delegateKey, beforeRewards := range beforeDelegateRewards {
+		afterRewards, ok := afterDelegateRewards[delegateKey]
+		require.True(t, ok)
+		_, _ = beforeRewards, afterRewards
+		// assert.EqualValues(t, beforeRewards.String(), afterRewards.String(), delegateKey)
+	}
+}
+
 func checkIBCTransferModule(t *testing.T, ctx sdk.Context, myApp *app.App) {
 	t.Helper()
 	escrowDenomsMap := nextversion.GetMigrateEscrowDenoms(ctx.ChainID())

app/upgrades/v8/distribution.go

@@ -7,12 +7,14 @@ import (
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	distributionkeeper "github.com/cosmos/cosmos-sdk/x/distribution/keeper"
 	"github.com/cosmos/cosmos-sdk/x/distribution/types"
+	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
 	"github.com/stretchr/testify/require"
 
 	fxtypes "github.com/pundiai/fx-core/v8/types"
+	fxstakingkeeper "github.com/pundiai/fx-core/v8/x/staking/keeper"
 )
 
-func migrateDistribution(ctx sdk.Context, distrKeeper distributionkeeper.Keeper) error {
+func migrateDistribution(ctx sdk.Context, stakingKeeper *fxstakingkeeper.Keeper, distrKeeper distributionkeeper.Keeper) error {
 	if err := migrateValidatorAccumulatedCommission(ctx, distrKeeper); err != nil {
 		return err
 	}
@@ -22,7 +24,7 @@ func migrateDistribution(ctx sdk.Context, distrKeeper distributionkeeper.Keeper)
 	if err := migrateValidatorCurrentRewards(ctx, distrKeeper); err != nil {
 		return err
 	}
-	if err := migrateDelegatorStartingInfos(ctx, distrKeeper); err != nil {
+	if err := migrateDelegatorStartingInfos(ctx, stakingKeeper, distrKeeper); err != nil {
 		return err
 	}
 	if err := migrateValidatorHistoricalRewards(ctx, distrKeeper); err != nil {
@@ -34,18 +36,39 @@ func migrateDistribution(ctx sdk.Context, distrKeeper distributionkeeper.Keeper)
 func migrateValidatorHistoricalRewards(ctx sdk.Context, keeper distributionkeeper.Keeper) error {
 	var err error
 	keeper.IterateValidatorHistoricalRewards(ctx, func(valAddr sdk.ValAddress, period uint64, rewards types.ValidatorHistoricalRewards) bool {
-		newRewards := swapFXDecCoinsToDefaultCoin(rewards.CumulativeRewardRatio)
+		newRewards := make(sdk.DecCoins, 0, len(rewards.CumulativeRewardRatio))
+		for _, coin := range rewards.CumulativeRewardRatio {
+			newDenom := coin.Denom
+			if newDenom == fxtypes.LegacyFXDenom {
+				newDenom = fxtypes.DefaultDenom
+			}
+			newRewards = append(newRewards, sdk.NewDecCoinFromDec(newDenom, coin.Amount))
+		}
 		rewards.CumulativeRewardRatio = newRewards
 		err = keeper.SetValidatorHistoricalRewards(ctx, valAddr, period, rewards)
 		return err != nil
 	})
 	return err
 }
 
-func migrateDelegatorStartingInfos(ctx sdk.Context, distrKeeper distributionkeeper.Keeper) error {
+func migrateDelegatorStartingInfos(ctx sdk.Context, stakingKeeper *fxstakingkeeper.Keeper, distrKeeper distributionkeeper.Keeper) error {
 	var err error
+	validatorMap := make(map[string]stakingtypes.Validator)
 	distrKeeper.IterateDelegatorStartingInfos(ctx, func(valAddr sdk.ValAddress, addr sdk.AccAddress, info types.DelegatorStartingInfo) bool {
-		info.Stake = fxtypes.SwapDecAmount(info.Stake)
+		var delegation stakingtypes.Delegation
+		delegation, err = stakingKeeper.GetDelegation(ctx, addr, valAddr)
+		if err != nil {
+			return true
+		}
+		validator, ok := validatorMap[delegation.ValidatorAddress]
+		if !ok {
+			validator, err = stakingKeeper.GetValidator(ctx, valAddr)
+			if err != nil {
+				return true
+			}
+			validatorMap[delegation.ValidatorAddress] = validator
+		}
+		info.Stake = validator.TokensFromSharesTruncated(delegation.Shares)
 		err = distrKeeper.SetDelegatorStartingInfo(ctx, valAddr, addr, info)
 		return err != nil
 	})
@@ -58,14 +81,14 @@ func migrateFeePool(ctx sdk.Context, distrKeeper distributionkeeper.Keeper) erro
 		return err
 	}
 
-	feePool.CommunityPool = swapFXDecCoinsToDefaultCoin(feePool.CommunityPool)
+	feePool.CommunityPool = fxtypes.SwapDecCoins(feePool.CommunityPool)
 	return distrKeeper.FeePool.Set(ctx, feePool)
 }
 
 func migrateValidatorAccumulatedCommission(ctx sdk.Context, distrKeeper distributionkeeper.Keeper) error {
 	var err error
 	distrKeeper.IterateValidatorAccumulatedCommissions(ctx, func(addr sdk.ValAddress, commission types.ValidatorAccumulatedCommission) bool {
-		newCommission := swapFXDecCoinsToDefaultCoin(commission.Commission)
+		newCommission := fxtypes.SwapDecCoins(commission.Commission)
 		commission.Commission = newCommission
 		err = distrKeeper.SetValidatorAccumulatedCommission(ctx, addr, commission)
 		return err != nil
@@ -77,7 +100,7 @@ func migrateValidatorAccumulatedCommission(ctx sdk.Context, distrKeeper distribu
 func migrateValidatorCurrentRewards(ctx sdk.Context, distrKeeper distributionkeeper.Keeper) error {
 	var err error
 	distrKeeper.IterateValidatorCurrentRewards(ctx, func(addr sdk.ValAddress, rewards types.ValidatorCurrentRewards) bool {
-		newRewards := swapFXDecCoinsToDefaultCoin(rewards.Rewards)
+		newRewards := fxtypes.SwapDecCoins(rewards.Rewards)
 		rewards.Rewards = newRewards
 		err = distrKeeper.SetValidatorCurrentRewards(ctx, addr, rewards)
 		return err != nil
@@ -88,7 +111,7 @@ func migrateValidatorCurrentRewards(ctx sdk.Context, distrKeeper distributionkee
 func migrateValidatorOutstandingRewards(ctx sdk.Context, distrKeeper distributionkeeper.Keeper) error {
 	var err error
 	distrKeeper.IterateValidatorOutstandingRewards(ctx, func(addr sdk.ValAddress, rewards types.ValidatorOutstandingRewards) bool {
-		newRewards := swapFXDecCoinsToDefaultCoin(rewards.Rewards)
+		newRewards := fxtypes.SwapDecCoins(rewards.Rewards)
 		rewards.Rewards = newRewards
 		err = distrKeeper.SetValidatorOutstandingRewards(ctx, addr, rewards)
 		return err != nil
@@ -97,23 +120,6 @@ func migrateValidatorOutstandingRewards(ctx sdk.Context, distrKeeper distributio
 	return err
 }
 
-func swapFXDecCoinsToDefaultCoin(decCoins sdk.DecCoins) sdk.DecCoins {
-	oldDenom := fxtypes.LegacyFXDenom
-	newDenom := fxtypes.DefaultDenom
-	newDecCoins := sdk.NewDecCoins()
-	for _, decCoin := range decCoins {
-		if decCoin.Denom == oldDenom {
-			decCoin.Denom = newDenom
-			decCoin.Amount = fxtypes.SwapDecAmount(decCoin.Amount)
-			if !decCoin.IsPositive() {
-				continue
-			}
-		}
-		newDecCoins = newDecCoins.Add(decCoin)
-	}
-	return newDecCoins
-}
-
 func CheckDistributionModule(t *testing.T, ctx sdk.Context, distrKeeper distributionkeeper.Keeper) {
 	t.Helper()
 

app/upgrades/v8/upgrade.go

@@ -228,7 +228,7 @@ func migrateModulesData(ctx sdk.Context, codec codec.Codec, app *keepers.AppKeep
 
 	migrateTransferTokenInEscrow(ctx, app.IBCTransferKeeper)
 
-	if err := migrateDistribution(ctx, app.DistrKeeper); err != nil {
+	if err := migrateDistribution(ctx, app.StakingKeeper, app.DistrKeeper); err != nil {
 		return err
 	}
 

go.mod

@@ -73,7 +73,7 @@ require (
 	github.com/iancoleman/strcase v0.3.0 // indirect
 	github.com/oasisprotocol/curve25519-voi v0.0.0-20230904125328-1f23a7beb09a // indirect
 	github.com/oklog/run v1.1.0 // indirect
-	github.com/rs/zerolog v1.33.0 // indirect
+	github.com/rs/zerolog v1.33.0
 	github.com/tyler-smith/go-bip39 v1.1.0 // indirect
 	gotest.tools/v3 v3.5.1 // indirect
 )

types/constant.go

@@ -118,3 +118,21 @@ func SwapCoins(coins sdk.Coins) sdk.Coins {
 	}
 	return swapCoins
 }
+
+func SwapDecCoins(decCoins sdk.DecCoins) sdk.DecCoins {
+	result := sdk.NewDecCoins()
+	if len(decCoins) == 0 {
+		return result
+	}
+	for _, decCoin := range decCoins {
+		if decCoin.Denom == LegacyFXDenom {
+			decCoin.Denom = DefaultDenom
+			decCoin.Amount = SwapDecAmount(decCoin.Amount)
+			if !decCoin.IsPositive() {
+				continue
+			}
+		}
+		result = result.Add(decCoin)
+	}
+	return result
+}