!test: add genesis state for multi staking module to simapp (#99)

* add genesis state for multi staking module to simapp * fix lint * Fix bonded pool balance and module balance * Move out const from module to testutil * Change variable names * update setup test * remove redundant file --------- Co-authored-by: Hieu Vu <[email protected]>
realiotech · Jan 13, 2024 · d96beeb · d96beeb
1 parent de5cbe2
commit d96beeb
Show file tree

Hide file tree

Showing 2 changed files with 103 additions and 37 deletions.
diff --git a/testing/simapp/test_helpers.go b/testing/simapp/test_helpers.go
@@ -9,6 +9,7 @@ import (
 	"testing"
 	"time"
 
+	multistakingtypes "github.com/realio-tech/multi-staking-module/x/multi-staking/types"
 	"github.com/stretchr/testify/require"
 	abci "github.com/tendermint/tendermint/abci/types"
 	"github.com/tendermint/tendermint/libs/log"
@@ -68,33 +69,19 @@ func setup(withGenesis bool, invCheckPeriod uint) (*SimApp, GenesisState) {
 
 // Setup initializes a new SimApp. A Nop logger is set in SimApp.
 func Setup(isCheckTx bool) *SimApp {
-	privVal := mock.NewPV()
-	pubKey, _ := privVal.GetPubKey()
-
-	// create validator set with single validator
-	validator := tmtypes.NewValidator(pubKey, 1)
-	valSet := tmtypes.NewValidatorSet([]*tmtypes.Validator{validator})
-
-	// generate genesis account
-	senderPrivKey := secp256k1.GenPrivKey()
-	acc := authtypes.NewBaseAccount(senderPrivKey.PubKey().Address().Bytes(), senderPrivKey.PubKey(), 0, 0)
-	balance := banktypes.Balance{
-		Address: acc.GetAddress().String(),
-		Coins:   sdk.NewCoins(sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(100000000000000))),
-	}
-
-	app := SetupWithGenesisValSet(valSet, []authtypes.GenesisAccount{acc}, balance)
+	valSet := GenValSet()
 
+	app := SetupWithGenesisValSet(valSet)
 	return app
 }
 
 // SetupWithGenesisValSet initializes a new SimApp with a validator set and genesis accounts
 // that also act as delegators. For simplicity, each validator is bonded with a delegation
 // of one consensus engine unit in the default token of the simapp from first genesis
 // account. A Nop logger is set in SimApp.
-func SetupWithGenesisValSet(valSet *tmtypes.ValidatorSet, genAccs []authtypes.GenesisAccount, balances ...banktypes.Balance) *SimApp {
+func SetupWithGenesisValSet(valSet *tmtypes.ValidatorSet) *SimApp {
 	app, genesisState := setup(true, 5)
-	genesisState = genesisStateWithValSet(app, genesisState, valSet, genAccs, balances...)
+	genesisState = genesisStateWithValSet(app, genesisState, valSet)
 
 	stateBytes, _ := json.MarshalIndent(genesisState, "", " ")
 
@@ -119,59 +106,100 @@ func SetupWithGenesisValSet(valSet *tmtypes.ValidatorSet, genAccs []authtypes.Ge
 	return app
 }
 
-func genesisStateWithValSet(app *SimApp, genesisState GenesisState,
-	valSet *tmtypes.ValidatorSet, genAccs []authtypes.GenesisAccount,
-	balances ...banktypes.Balance,
-) GenesisState {
+func genesisStateWithValSet(app *SimApp, genesisState GenesisState, valSet *tmtypes.ValidatorSet) GenesisState {
+	genAcc := GenAcc()
+	genAccs := []authtypes.GenesisAccount{genAcc}
+	balances := []banktypes.Balance{}
+
 	// set genesis accounts
 	authGenesis := authtypes.NewGenesisState(authtypes.DefaultParams(), genAccs)
 	genesisState[authtypes.ModuleName] = app.AppCodec().MustMarshalJSON(authGenesis)
 
+	// set multi staking genesis state
+	msCoinAInfo := multistakingtypes.MultiStakingCoinInfo{
+		Denom:      MultiStakingCoinA.Denom,
+		BondWeight: MultiStakingCoinA.BondWeight,
+	}
+	msCoinBInfo := multistakingtypes.MultiStakingCoinInfo{
+		Denom:      MultiStakingCoinB.Denom,
+		BondWeight: MultiStakingCoinB.BondWeight,
+	}
+	msCoinInfos := []multistakingtypes.MultiStakingCoinInfo{msCoinAInfo, msCoinBInfo}
+	validatorMsCoins := make([]multistakingtypes.ValidatorMultiStakingCoin, 0, len(valSet.Validators))
+	locks := make([]multistakingtypes.MultiStakingLock, 0, len(valSet.Validators))
+	lockCoins := sdk.NewCoins()
+
+	// staking genesis state
 	validators := make([]stakingtypes.Validator, 0, len(valSet.Validators))
 	delegations := make([]stakingtypes.Delegation, 0, len(valSet.Validators))
+	bondCoins := sdk.NewCoins()
 
-	bondAmt := sdk.DefaultPowerReduction
+	for i, val := range valSet.Validators {
+		valMsCoin := MultiStakingCoinA
+		if i%2 == 1 {
+			valMsCoin = MultiStakingCoinB
+		}
+
+		validatorMsCoins = append(validatorMsCoins, multistakingtypes.ValidatorMultiStakingCoin{
+			ValAddr:   sdk.ValAddress(val.Address).String(),
+			CoinDenom: valMsCoin.Denom,
+		})
+
+		lockId := multistakingtypes.MultiStakingLockID(genAcc.GetAddress().String(), sdk.ValAddress(val.Address).String())
+		lockRecord := multistakingtypes.NewMultiStakingLock(&lockId, valMsCoin)
+
+		locks = append(locks, lockRecord)
+		lockCoins = lockCoins.Add(valMsCoin.ToCoin())
 
-	for _, val := range valSet.Validators {
 		pk, _ := cryptocodec.FromTmPubKeyInterface(val.PubKey)
 		pkAny, _ := codectypes.NewAnyWithValue(pk)
 		validator := stakingtypes.Validator{
 			OperatorAddress:   sdk.ValAddress(val.Address).String(),
 			ConsensusPubkey:   pkAny,
 			Jailed:            false,
 			Status:            stakingtypes.Bonded,
-			Tokens:            bondAmt,
+			Tokens:            valMsCoin.BondValue(),
 			DelegatorShares:   sdk.OneDec(),
 			Description:       stakingtypes.Description{},
 			UnbondingHeight:   int64(0),
 			UnbondingTime:     time.Unix(0, 0).UTC(),
 			Commission:        stakingtypes.NewCommission(sdk.ZeroDec(), sdk.ZeroDec(), sdk.ZeroDec()),
 			MinSelfDelegation: sdk.ZeroInt(),
 		}
+
 		validators = append(validators, validator)
-		delegations = append(delegations, stakingtypes.NewDelegation(genAccs[0].GetAddress(), val.Address.Bytes(), sdk.OneDec()))
+		delegations = append(delegations, stakingtypes.NewDelegation(genAcc.GetAddress(), val.Address.Bytes(), sdk.OneDec()))
 
+		bondCoins = bondCoins.Add(sdk.NewCoin(sdk.DefaultBondDenom, valMsCoin.BondValue()))
 	}
+
 	// set validators and delegations
 	stakingGenesis := stakingtypes.NewGenesisState(stakingtypes.DefaultParams(), validators, delegations)
 	genesisState[stakingtypes.ModuleName] = app.AppCodec().MustMarshalJSON(stakingGenesis)
 
-	totalSupply := sdk.NewCoins()
-	for _, b := range balances {
-		// add genesis acc tokens to total supply
-		totalSupply = totalSupply.Add(b.Coins...)
-	}
-
-	for range delegations {
-		// add delegated tokens to total supply
-		totalSupply = totalSupply.Add(sdk.NewCoin(sdk.DefaultBondDenom, bondAmt))
+	multistakingGenesis := multistakingtypes.GenesisState{
+		MultiStakingLocks:          locks,
+		MultiStakingUnlocks:        []multistakingtypes.MultiStakingUnlock{},
+		MultiStakingCoinInfo:       msCoinInfos,
+		ValidatorMultiStakingCoins: validatorMsCoins,
+		StakingGenesisState:        stakingGenesis,
 	}
+	genesisState[multistakingtypes.ModuleName] = app.AppCodec().MustMarshalJSON(&multistakingGenesis)
 
-	// add bonded amount to bonded pool module account
 	balances = append(balances, banktypes.Balance{
 		Address: authtypes.NewModuleAddress(stakingtypes.BondedPoolName).String(),
-		Coins:   sdk.Coins{sdk.NewCoin(sdk.DefaultBondDenom, bondAmt)},
+		Coins:   bondCoins,
 	})
+	balances = append(balances, banktypes.Balance{
+		Address: authtypes.NewModuleAddress(multistakingtypes.ModuleName).String(),
+		Coins:   lockCoins,
+	})
+
+	totalSupply := sdk.NewCoins()
+	for _, b := range balances {
+		// add genesis acc tokens to total supply
+		totalSupply = totalSupply.Add(b.Coins...)
+	}
 
 	// update total supply
 	bankGenesis := banktypes.NewGenesisState(banktypes.DefaultGenesisState().Params, balances, totalSupply, []banktypes.Metadata{})
@@ -459,3 +487,24 @@ func FundAccount(app *SimApp, ctx sdk.Context, addr sdk.AccAddress, amounts sdk.
 	}
 	return app.BankKeeper.SendCoinsFromModuleToAccount(ctx, minttypes.ModuleName, addr, amounts)
 }
+
+func GenValSet() *tmtypes.ValidatorSet {
+	privVal0 := mock.NewPV()
+	privVal1 := mock.NewPV()
+
+	pubKey0, _ := privVal0.GetPubKey()
+	pubKey1, _ := privVal1.GetPubKey()
+
+	// create validator set with single validator
+	val0 := tmtypes.NewValidator(pubKey0, 1)
+	val1 := tmtypes.NewValidator(pubKey1, 1)
+
+	valSet := tmtypes.NewValidatorSet([]*tmtypes.Validator{val0, val1})
+
+	return valSet
+}
+
+func GenAcc() authtypes.GenesisAccount {
+	senderPrivKey := secp256k1.GenPrivKey()
+	return authtypes.NewBaseAccount(senderPrivKey.PubKey().Address().Bytes(), senderPrivKey.PubKey(), 0, 0)
+}
diff --git a/testing/simapp/types.go b/testing/simapp/types.go
@@ -1,6 +1,7 @@
 package simapp
 
 import (
+	multistakingtypes "github.com/realio-tech/multi-staking-module/x/multi-staking/types"
 	abci "github.com/tendermint/tendermint/abci/types"
 
 	"github.com/cosmos/cosmos-sdk/codec"
@@ -9,6 +10,22 @@ import (
 	"github.com/cosmos/cosmos-sdk/types/module"
 )
 
+var (
+	MultiStakingCoinA = multistakingtypes.MultiStakingCoin{
+		Denom:      "ario",
+		Amount:     sdk.NewIntFromUint64(100000000),
+		BondWeight: sdk.MustNewDecFromStr("1.23"),
+	}
+	MultiStakingCoinB = multistakingtypes.MultiStakingCoin{
+		Denom:      "arst",
+		Amount:     sdk.NewIntFromUint64(100000000),
+		BondWeight: sdk.MustNewDecFromStr("0.12"),
+	}
+	MultiStakingCoins = []multistakingtypes.MultiStakingCoin{
+		MultiStakingCoinA, MultiStakingCoinB,
+	}
+)
+
 // App implements the common methods for a Cosmos SDK-based application
 // specific blockchain.
 type App interface {