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

     1  package v4_test
     2  
     3  import (
     4  	"testing"
     5  	"time"
     6  
     7  	"github.com/stretchr/testify/require"
     8  
     9  	"cosmossdk.io/math"
    10  	storetypes "cosmossdk.io/store/types"
    11  
    12  	"github.com/cosmos/cosmos-sdk/codec"
    13  	"github.com/cosmos/cosmos-sdk/testutil"
    14  	"github.com/cosmos/cosmos-sdk/testutil/sims"
    15  	sdk "github.com/cosmos/cosmos-sdk/types"
    16  	moduletestutil "github.com/cosmos/cosmos-sdk/types/module/testutil"
    17  	paramtypes "github.com/cosmos/cosmos-sdk/x/params/types"
    18  	"github.com/cosmos/cosmos-sdk/x/staking"
    19  	v4 "github.com/cosmos/cosmos-sdk/x/staking/migrations/v4"
    20  	"github.com/cosmos/cosmos-sdk/x/staking/types"
    21  )
    22  
    23  type mockSubspace struct {
    24  	ps types.Params
    25  }
    26  
    27  func newMockSubspace(ps types.Params) mockSubspace {
    28  	return mockSubspace{ps: ps}
    29  }
    30  
    31  func (ms mockSubspace) GetParamSet(ctx sdk.Context, ps paramtypes.ParamSet) {
    32  	*ps.(*types.Params) = ms.ps
    33  }
    34  
    35  func TestMigrate(t *testing.T) {
    36  	cdc := moduletestutil.MakeTestEncodingConfig(staking.AppModuleBasic{}).Codec
    37  
    38  	storeKey := storetypes.NewKVStoreKey(v4.ModuleName)
    39  	tKey := storetypes.NewTransientStoreKey("transient_test")
    40  	ctx := testutil.DefaultContext(storeKey, tKey)
    41  	store := ctx.KVStore(storeKey)
    42  	duplicateCreationHeight := int64(1)
    43  
    44  	accAddrs := sims.CreateIncrementalAccounts(1)
    45  	accAddr := accAddrs[0]
    46  
    47  	valAddrs := sims.ConvertAddrsToValAddrs(accAddrs)
    48  	valAddr := valAddrs[0]
    49  
    50  	// creating 10 ubdEntries with same height and 10 ubdEntries with different creation height
    51  	err := createOldStateUnbonding(t, duplicateCreationHeight, valAddr, accAddr, cdc, store)
    52  	require.NoError(t, err)
    53  
    54  	legacySubspace := newMockSubspace(types.DefaultParams())
    55  
    56  	testCases := []struct {
    57  		name        string
    58  		doMigration bool
    59  	}{
    60  		{
    61  			name:        "without state migration",
    62  			doMigration: false,
    63  		},
    64  		{
    65  			name:        "with state migration",
    66  			doMigration: true,
    67  		},
    68  	}
    69  
    70  	for _, tc := range testCases {
    71  		t.Run(tc.name, func(t *testing.T) {
    72  			if tc.doMigration {
    73  				require.NoError(t, v4.MigrateStore(ctx, store, cdc, legacySubspace))
    74  			}
    75  
    76  			ubd := getUBD(t, accAddr, valAddr, store, cdc)
    77  			if tc.doMigration {
    78  				var res types.Params
    79  				bz := store.Get(v4.ParamsKey)
    80  				require.NoError(t, cdc.Unmarshal(bz, &res))
    81  				require.Equal(t, legacySubspace.ps, res)
    82  
    83  				// checking the updated balance for duplicateCreationHeight
    84  				for _, ubdEntry := range ubd.Entries {
    85  					if ubdEntry.CreationHeight == duplicateCreationHeight {
    86  						require.Equal(t, math.NewInt(100*10), ubdEntry.Balance)
    87  						break
    88  					}
    89  				}
    90  				require.Equal(t, 11, len(ubd.Entries))
    91  			} else {
    92  				require.Equal(t, true, true)
    93  				require.Equal(t, 20, len(ubd.Entries))
    94  			}
    95  		})
    96  	}
    97  }
    98  
    99  // createOldStateUnbonding will create the ubd entries with duplicate heights
   100  // 10 duplicate heights and 10 unique ubd with creation height
   101  func createOldStateUnbonding(t *testing.T, creationHeight int64, valAddr sdk.ValAddress, accAddr sdk.AccAddress, cdc codec.BinaryCodec, store storetypes.KVStore) error {
   102  	unbondBalance := math.NewInt(100)
   103  	completionTime := time.Now()
   104  	ubdEntries := make([]types.UnbondingDelegationEntry, 0, 10)
   105  
   106  	for i := int64(0); i < 10; i++ {
   107  		ubdEntry := types.UnbondingDelegationEntry{
   108  			CreationHeight: creationHeight,
   109  			Balance:        unbondBalance,
   110  			InitialBalance: unbondBalance,
   111  			CompletionTime: completionTime,
   112  		}
   113  		ubdEntries = append(ubdEntries, ubdEntry)
   114  		// creating more entries for testing the creation_heights
   115  		ubdEntry.CreationHeight = i + 2
   116  		ubdEntry.CompletionTime = completionTime.Add(time.Minute * 10)
   117  		ubdEntries = append(ubdEntries, ubdEntry)
   118  	}
   119  
   120  	ubd := types.UnbondingDelegation{
   121  		ValidatorAddress: valAddr.String(),
   122  		DelegatorAddress: accAddr.String(),
   123  		Entries:          ubdEntries,
   124  	}
   125  
   126  	// set the unbond delegation with validator and delegator
   127  	bz := types.MustMarshalUBD(cdc, ubd)
   128  	key := getUBDKey(accAddr, valAddr)
   129  	store.Set(key, bz)
   130  	return nil
   131  }
   132  
   133  func getUBD(t *testing.T, accAddr sdk.AccAddress, valAddr sdk.ValAddress, store storetypes.KVStore, cdc codec.BinaryCodec) types.UnbondingDelegation {
   134  	// get the unbonding delegations
   135  	var ubdRes types.UnbondingDelegation
   136  	ubdbz := store.Get(getUBDKey(accAddr, valAddr))
   137  	require.NoError(t, cdc.Unmarshal(ubdbz, &ubdRes))
   138  	return ubdRes
   139  }
   140  
   141  func getUBDKey(accAddr sdk.AccAddress, valAddr sdk.ValAddress) []byte {
   142  	return types.GetUBDKey(accAddr, valAddr)
   143  }