github.com/cosmos/cosmos-sdk@v0.50.10/x/staking/migrations/v5/migrations_test.go (about)

     1  package v5_test
     2  
     3  import (
     4  	"math"
     5  	"math/rand"
     6  	"testing"
     7  	"time"
     8  
     9  	cmtproto "github.com/cometbft/cometbft/proto/tendermint/types"
    10  	"github.com/stretchr/testify/assert"
    11  	"github.com/stretchr/testify/require"
    12  
    13  	sdkmath "cosmossdk.io/math"
    14  	storetypes "cosmossdk.io/store/types"
    15  
    16  	"github.com/cosmos/cosmos-sdk/codec"
    17  	"github.com/cosmos/cosmos-sdk/testutil"
    18  	"github.com/cosmos/cosmos-sdk/testutil/sims"
    19  	sdk "github.com/cosmos/cosmos-sdk/types"
    20  	moduletestutil "github.com/cosmos/cosmos-sdk/types/module/testutil"
    21  	"github.com/cosmos/cosmos-sdk/x/staking"
    22  	v2 "github.com/cosmos/cosmos-sdk/x/staking/migrations/v2"
    23  	v5 "github.com/cosmos/cosmos-sdk/x/staking/migrations/v5"
    24  	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
    25  )
    26  
    27  func TestHistoricalKeysMigration(t *testing.T) {
    28  	storeKey := storetypes.NewKVStoreKey("staking")
    29  	tKey := storetypes.NewTransientStoreKey("transient_test")
    30  	ctx := testutil.DefaultContext(storeKey, tKey)
    31  	store := ctx.KVStore(storeKey)
    32  
    33  	type testCase struct {
    34  		oldKey, newKey []byte
    35  		historicalInfo []byte
    36  	}
    37  
    38  	testCases := make(map[int64]testCase)
    39  
    40  	// edge cases
    41  	testCases[0], testCases[1], testCases[math.MaxInt32] = testCase{}, testCase{}, testCase{}
    42  
    43  	// random cases
    44  	seed := time.Now().UnixNano()
    45  	r := rand.New(rand.NewSource(seed))
    46  	for i := 0; i < 10; i++ {
    47  		height := r.Intn(math.MaxInt32-2) + 2
    48  
    49  		testCases[int64(height)] = testCase{}
    50  	}
    51  
    52  	cdc := moduletestutil.MakeTestEncodingConfig().Codec
    53  	for height := range testCases {
    54  		testCases[height] = testCase{
    55  			oldKey:         v2.GetHistoricalInfoKey(height),
    56  			newKey:         v5.GetHistoricalInfoKey(height),
    57  			historicalInfo: cdc.MustMarshal(createHistoricalInfo(height, "testChainID")),
    58  		}
    59  	}
    60  
    61  	// populate store using old key format
    62  	for _, tc := range testCases {
    63  		store.Set(tc.oldKey, tc.historicalInfo)
    64  	}
    65  
    66  	// migrate store to new key format
    67  	require.NoErrorf(t, v5.MigrateStore(ctx, store, cdc), "v5.MigrateStore failed, seed: %d", seed)
    68  
    69  	// check results
    70  	for _, tc := range testCases {
    71  		require.Nilf(t, store.Get(tc.oldKey), "old key should be deleted, seed: %d", seed)
    72  		require.NotNilf(t, store.Get(tc.newKey), "new key should be created, seed: %d", seed)
    73  		require.Equalf(t, tc.historicalInfo, store.Get(tc.newKey), "seed: %d", seed)
    74  	}
    75  }
    76  
    77  func createHistoricalInfo(height int64, chainID string) *stakingtypes.HistoricalInfo {
    78  	return &stakingtypes.HistoricalInfo{Header: cmtproto.Header{ChainID: chainID, Height: height}}
    79  }
    80  
    81  func TestDelegationsByValidatorMigrations(t *testing.T) {
    82  	cdc := moduletestutil.MakeTestEncodingConfig(staking.AppModuleBasic{}).Codec
    83  	storeKey := storetypes.NewKVStoreKey(v5.ModuleName)
    84  	tKey := storetypes.NewTransientStoreKey("transient_test")
    85  	ctx := testutil.DefaultContext(storeKey, tKey)
    86  	store := ctx.KVStore(storeKey)
    87  
    88  	accAddrs := sims.CreateIncrementalAccounts(11)
    89  	valAddrs := sims.ConvertAddrsToValAddrs(accAddrs[0:1])
    90  	var addedDels []stakingtypes.Delegation
    91  
    92  	for i := 1; i < 11; i++ {
    93  		del1 := stakingtypes.NewDelegation(accAddrs[i].String(), valAddrs[0].String(), sdkmath.LegacyNewDec(100))
    94  		store.Set(stakingtypes.GetDelegationKey(accAddrs[i], valAddrs[0]), stakingtypes.MustMarshalDelegation(cdc, del1))
    95  		addedDels = append(addedDels, del1)
    96  	}
    97  
    98  	// before migration the state of delegations by val index should be empty
    99  	dels := getValDelegations(ctx, cdc, storeKey, valAddrs[0])
   100  	assert.Len(t, dels, 0)
   101  
   102  	err := v5.MigrateStore(ctx, store, cdc)
   103  	assert.NoError(t, err)
   104  
   105  	// after migration the state of delegations by val index should not be empty
   106  	dels = getValDelegations(ctx, cdc, storeKey, valAddrs[0])
   107  	assert.Len(t, dels, len(addedDels))
   108  	assert.Equal(t, addedDels, dels)
   109  }
   110  
   111  func getValDelegations(ctx sdk.Context, cdc codec.Codec, storeKey storetypes.StoreKey, valAddr sdk.ValAddress) []stakingtypes.Delegation {
   112  	var delegations []stakingtypes.Delegation
   113  
   114  	store := ctx.KVStore(storeKey)
   115  	iterator := storetypes.KVStorePrefixIterator(store, v5.GetDelegationsByValPrefixKey(valAddr))
   116  	for ; iterator.Valid(); iterator.Next() {
   117  		var delegation stakingtypes.Delegation
   118  		valAddr, delAddr, err := stakingtypes.ParseDelegationsByValKey(iterator.Key())
   119  		if err != nil {
   120  			panic(err)
   121  		}
   122  
   123  		bz := store.Get(stakingtypes.GetDelegationKey(delAddr, valAddr))
   124  
   125  		cdc.MustUnmarshal(bz, &delegation)
   126  
   127  		delegations = append(delegations, delegation)
   128  	}
   129  
   130  	return delegations
   131  }