chore(tests): add restaking simulation tests (#212)

## Description

Closes: MILK-150

<!-- Add a description of the changes that this PR introduces and the
files that
are the most critical to review. -->

---

### Author Checklist

*All items are required. Please add a note to the item if the item is
not applicable and
please add links to any relevant follow up issues.*

I have...

- [x] included the correct [type
prefix](https://github.com/commitizen/conventional-commit-types/blob/v3.0.0/index.json)
in the PR title
- [ ] added `!` to the type prefix if API or client breaking change
- [ ] targeted the correct branch (see [PR
Targeting](https://github.com/milkyway-labs/milkyway/blob/master/CONTRIBUTING.md#pr-targeting))
- [ ] provided a link to the relevant issue or specification
- [ ] followed the guidelines for [building
modules](https://docs.cosmos.network/v0.44/building-modules/intro.html)
- [ ] included the necessary unit and integration
[tests](https://github.com/milkyway-labs/milkyway/blob/master/CONTRIBUTING.md#testing)
- [ ] added a changelog entry to `CHANGELOG.md`
- [ ] included comments for [documenting Go
code](https://blog.golang.org/godoc)
- [ ] updated the relevant documentation or specification
- [x] reviewed "Files changed" and left comments if necessary
- [x] confirmed all CI checks have passed

### Reviewers Checklist

*All items are required. Please add a note if the item is not applicable
and please add
your handle next to the items reviewed if you only reviewed selected
items.*

I have...

- [ ] confirmed the correct [type
prefix](https://github.com/commitizen/conventional-commit-types/blob/v3.0.0/index.json)
in the PR title
- [ ] confirmed `!` in the type prefix if API or client breaking change
- [ ] confirmed all author checklist items have been addressed
- [ ] reviewed state machine logic
- [ ] reviewed API design and naming
- [ ] reviewed documentation is accurate
- [ ] reviewed tests and test coverage
- [ ] manually tested (if applicable)

app/modules.go

@@ -152,7 +152,7 @@ func appModules(
 		services.NewAppModule(appCodec, app.ServicesKeeper, app.AccountKeeper, app.BankKeeper),
 		operators.NewAppModule(appCodec, app.OperatorsKeeper, app.AccountKeeper, app.BankKeeper, app.StakingKeeper),
 		pools.NewAppModule(appCodec, app.PoolsKeeper),
-		restaking.NewAppModule(appCodec, app.RestakingKeeper),
+		restaking.NewAppModule(appCodec, app.RestakingKeeper, app.AccountKeeper, app.BankKeeper, app.PoolsKeeper, app.OperatorsKeeper, app.ServicesKeeper),
 		assets.NewAppModule(appCodec, app.AssetsKeeper),
 		rewards.NewAppModule(appCodec, app.RewardsKeeper),
 		liquidvesting.NewAppModule(appCodec, app.LiquidVestingKeeper),
@@ -205,6 +205,7 @@ func simulationModules(
 		// MilkyWay modules
 		services.NewAppModule(appCodec, app.ServicesKeeper, app.AccountKeeper, app.BankKeeper),
 		operators.NewAppModule(appCodec, app.OperatorsKeeper, app.AccountKeeper, app.BankKeeper, app.StakingKeeper),
+		restaking.NewAppModule(appCodec, app.RestakingKeeper, app.AccountKeeper, app.BankKeeper, app.PoolsKeeper, app.OperatorsKeeper, app.ServicesKeeper),
 	}
 }
 

testutils/simtesting/utils.go

@@ -87,3 +87,33 @@ func RandomCoin(r *rand.Rand, denom string, maxAmount int) sdk.Coin {
 		sdkmath.NewInt(int64(r.Intn(maxAmount*1e6))),
 	)
 }
+
+// RandomSubSlice returns a random subset of the given slice
+func RandomSubSlice[T any](r *rand.Rand, items []T) []T {
+	// Empty slice, we can't pick random elements
+	if len(items) == 0 {
+		return nil
+	}
+
+	// We store here the selected index, this allows T to not be comparable.
+	pickedIndexes := make(map[int]bool)
+
+	var elements []T
+	// Randomly select how many items to pick
+	count := r.Intn(len(items) + 1)
+	for len(pickedIndexes) < count {
+		// Get a random index
+		index := r.Intn(len(items))
+		_, found := pickedIndexes[index]
+
+		// Check if we have already picked this element
+		if !found {
+			// Element not picked, add it
+			elements = append(elements, items[index])
+			// Signal that we have picked the element at index
+			pickedIndexes[index] = true
+		}
+	}
+
+	return elements
+}

x/operators/simulation/genesis.go

@@ -58,3 +58,16 @@ func RandomizedGenState(simState *module.SimulationState) {
 	genesis := types.NewGenesisState(nextOperatorID, operators, operatorParams, unbondingOperators, params)
 	simState.GenState[types.ModuleName] = simState.Cdc.MustMarshalJSON(genesis)
 }
+
+// GetGenesisState returns the operators genesis state from the SimulationState
+func GetGenesisState(simState *module.SimulationState) types.GenesisState {
+	operatorsGenesisJSON, found := simState.GenState[types.ModuleName]
+	var operatorsGenesis types.GenesisState
+	if found {
+		simState.Cdc.MustUnmarshalJSON(operatorsGenesisJSON, &operatorsGenesis)
+	} else {
+		operatorsGenesis = *types.DefaultGenesis()
+	}
+
+	return operatorsGenesis
+}

x/pools/simulation/genesis.go

@@ -0,0 +1,20 @@
+package simulation
+
+import (
+	"github.com/cosmos/cosmos-sdk/types/module"
+
+	"github.com/milkyway-labs/milkyway/v3/x/pools/types"
+)
+
+// GetGenesisState returns the pools genesis state from the SimulationState
+func GetGenesisState(simState *module.SimulationState) types.GenesisState {
+	operatorsGenesisJSON, found := simState.GenState[types.ModuleName]
+	var operatorsGenesis types.GenesisState
+	if found {
+		simState.Cdc.MustUnmarshalJSON(operatorsGenesisJSON, &operatorsGenesis)
+	} else {
+		operatorsGenesis = *types.DefaultGenesis()
+	}
+
+	return operatorsGenesis
+}

x/pools/simulation/utils.go

@@ -0,0 +1,30 @@
+package simulation
+
+import (
+	"math/rand"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+
+	"github.com/milkyway-labs/milkyway/v3/utils"
+	"github.com/milkyway-labs/milkyway/v3/x/pools/keeper"
+	"github.com/milkyway-labs/milkyway/v3/x/pools/types"
+)
+
+// GetRandomExistingPool returns a random existing pool
+func GetRandomExistingPool(r *rand.Rand, ctx sdk.Context, k *keeper.Keeper, filter func(s types.Pool) bool) (types.Pool, bool) {
+	pools, err := k.GetPools(ctx)
+	if err != nil {
+		panic(err)
+	}
+
+	if filter != nil {
+		pools = utils.Filter(pools, filter)
+	}
+
+	if len(pools) == 0 {
+		return types.Pool{}, false
+	}
+
+	randomServiceIndex := r.Intn(len(pools))
+	return pools[randomServiceIndex], true
+}

x/restaking/keeper/invariants.go

@@ -373,7 +373,7 @@ func AllowedOperatorsExistInvariant(k *Keeper) sdk.Invariant {
 		// Iterate over all the services joined by operators
 		var notFoundOperatorsIDs []uint32
 		err := k.IterateAllServicesAllowedOperators(ctx, func(serviceID uint32, operatorID uint32) (stop bool, err error) {
-			_, err = k.operatorsKeeper.GetOperator(ctx, serviceID)
+			_, err = k.operatorsKeeper.GetOperator(ctx, operatorID)
 			if err != nil {
 				if errors.Is(err, collections.ErrNotFound) {
 					notFoundOperatorsIDs = append(notFoundOperatorsIDs, operatorID)

x/restaking/keeper/keeper.go

@@ -26,7 +26,7 @@ type Keeper struct {
 	servicesKeeper  types.ServicesKeeper
 
 	// Keeper data
-	schema collections.Schema
+	Schema collections.Schema
 
 	// Here we use a IndexMap with NoValue instead of a KeySet because the cosmos-sdk don't
 	// provide a KeySet with indexes that we need in order to get the list of operators
@@ -103,7 +103,7 @@ func NewKeeper(
 	if err != nil {
 		panic(err)
 	}
-	k.schema = schema
+	k.Schema = schema
 
 	return k
 }

x/restaking/module.go

@@ -6,20 +6,25 @@ import (
 	"fmt"
 
 	"cosmossdk.io/core/appmodule"
-	"github.com/grpc-ecosystem/grpc-gateway/runtime"
-	"github.com/spf13/cobra"
-
 	abci "github.com/cometbft/cometbft/abci/types"
-
 	"github.com/cosmos/cosmos-sdk/client"
 	"github.com/cosmos/cosmos-sdk/codec"
 	cdctypes "github.com/cosmos/cosmos-sdk/codec/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/module"
+	simtypes "github.com/cosmos/cosmos-sdk/types/simulation"
+	authkeeper "github.com/cosmos/cosmos-sdk/x/auth/keeper"
+	bankkeeper "github.com/cosmos/cosmos-sdk/x/bank/keeper"
+	"github.com/grpc-ecosystem/grpc-gateway/runtime"
+	"github.com/spf13/cobra"
 
+	operatorskeeper "github.com/milkyway-labs/milkyway/v3/x/operators/keeper"
+	poolskeeper "github.com/milkyway-labs/milkyway/v3/x/pools/keeper"
 	"github.com/milkyway-labs/milkyway/v3/x/restaking/client/cli"
 	"github.com/milkyway-labs/milkyway/v3/x/restaking/keeper"
+	"github.com/milkyway-labs/milkyway/v3/x/restaking/simulation"
 	"github.com/milkyway-labs/milkyway/v3/x/restaking/types"
+	serviceskeeper "github.com/milkyway-labs/milkyway/v3/x/services/keeper"
 )
 
 const (
@@ -100,12 +105,31 @@ type AppModule struct {
 
 	// To ensure setting hooks properly, keeper must be a reference
 	keeper *keeper.Keeper
-}
 
-func NewAppModule(cdc codec.Codec, keeper *keeper.Keeper) AppModule {
+	ak              authkeeper.AccountKeeper
+	bk              bankkeeper.Keeper
+	poolsKeeper     *poolskeeper.Keeper
+	operatorsKeeper *operatorskeeper.Keeper
+	servicesKeeper  *serviceskeeper.Keeper
+}
+
+func NewAppModule(
+	cdc codec.Codec,
+	keeper *keeper.Keeper,
+	ak authkeeper.AccountKeeper,
+	bk bankkeeper.Keeper,
+	poolsKeeper *poolskeeper.Keeper,
+	operatorsKeeper *operatorskeeper.Keeper,
+	servicesKeeper *serviceskeeper.Keeper,
+) AppModule {
 	return AppModule{
-		AppModuleBasic: NewAppModuleBasic(cdc),
-		keeper:         keeper,
+		AppModuleBasic:  NewAppModuleBasic(cdc),
+		keeper:          keeper,
+		ak:              ak,
+		bk:              bk,
+		poolsKeeper:     poolsKeeper,
+		operatorsKeeper: operatorsKeeper,
+		servicesKeeper:  servicesKeeper,
 	}
 }
 
@@ -156,3 +180,33 @@ func (am AppModule) EndBlock(ctx context.Context) error {
 func (am AppModule) IsOnePerModuleType() {}
 
 func (am AppModule) IsAppModule() {}
+
+// AppModuleSimulation functions
+
+// GenerateGenesisState creates a randomized GenState of the restaking module.
+func (AppModule) GenerateGenesisState(simState *module.SimulationState) {
+	simulation.RandomizedGenState(simState)
+}
+
+// ProposalMsgs returns msgs used for governance proposals for simulations.
+func (am AppModule) ProposalMsgs(simState module.SimulationState) []simtypes.WeightedProposalMsg {
+	return simulation.ProposalMsgs(am.keeper)
+}
+
+// RegisterStoreDecoder registers a decoder for restaking module's types
+func (am AppModule) RegisterStoreDecoder(sdr simtypes.StoreDecoderRegistry) {
+	sdr[types.StoreKey] = simulation.NewDecodeStore(am.cdc, am.keeper)
+}
+
+// WeightedOperations returns all the restaking module operations with their respective weights.
+func (am AppModule) WeightedOperations(simState module.SimulationState) []simtypes.WeightedOperation {
+	return simulation.WeightedOperations(
+		simState.AppParams,
+		am.ak,
+		am.bk,
+		am.poolsKeeper,
+		am.operatorsKeeper,
+		am.servicesKeeper,
+		am.keeper,
+	)
+}

x/restaking/simulation/decoder.go

@@ -0,0 +1,84 @@
+package simulation
+
+import (
+	"bytes"
+	"encoding/binary"
+	"fmt"
+
+	"github.com/cosmos/cosmos-sdk/codec"
+	"github.com/cosmos/cosmos-sdk/types/kv"
+	simtypes "github.com/cosmos/cosmos-sdk/types/simulation"
+
+	"github.com/milkyway-labs/milkyway/v3/x/restaking/keeper"
+	"github.com/milkyway-labs/milkyway/v3/x/restaking/types"
+)
+
+// NewDecodeStore returns a decoder function closure that unmarshals the KVPair's
+// Value to the corresponding restaking type.
+func NewDecodeStore(cdc codec.BinaryCodec, keeper *keeper.Keeper) func(kvA kv.Pair, kvB kv.Pair) string {
+	collectionsDecoder := simtypes.NewStoreDecoderFuncFromCollectionsSchema(keeper.Schema)
+
+	return func(kvA, kvB kv.Pair) string {
+		switch {
+
+		case bytes.Equal(kvA.Key[:1], types.UnbondingIDKey):
+			idA := binary.BigEndian.Uint64(kvA.Value)
+			idB := binary.BigEndian.Uint64(kvB.Value)
+			return fmt.Sprintf("%d\n%d", idA, idB)
+
+		case bytes.Equal(kvA.Key[:1], types.UnbondingIndexKey):
+			return fmt.Sprintf("%v\n%v", kvA.Value, kvB.Value)
+
+		case bytes.Equal(kvA.Key[:1], types.UnbondingTypeKey):
+			typeA := binary.BigEndian.Uint32(kvA.Value)
+			typeB := binary.BigEndian.Uint32(kvB.Value)
+			return fmt.Sprintf("%d\n%d", typeA, typeB)
+
+		case bytes.Equal(kvA.Key[:1], types.PoolDelegationPrefix):
+			delegationA := types.MustUnmarshalDelegation(cdc, kvA.Value)
+			delegationB := types.MustUnmarshalDelegation(cdc, kvB.Value)
+			return fmt.Sprintf("%v\n%v", delegationA, delegationB)
+
+		case bytes.Equal(kvA.Key[:1], types.PoolUnbondingDelegationPrefix):
+			undelegationA := types.MustUnmarshalUnbondingDelegation(cdc, kvA.Value)
+			undelegationB := types.MustUnmarshalUnbondingDelegation(cdc, kvB.Value)
+			return fmt.Sprintf("%v\n%v", undelegationA, undelegationB)
+
+		case bytes.Equal(kvA.Key[:1], types.OperatorDelegationPrefix):
+			delegationA := types.MustUnmarshalDelegation(cdc, kvA.Value)
+			delegationB := types.MustUnmarshalDelegation(cdc, kvB.Value)
+			return fmt.Sprintf("%v\n%v", delegationA, delegationB)
+
+		case bytes.Equal(kvA.Key[:1], types.OperatorUnbondingDelegationPrefix):
+			undelegationA := types.MustUnmarshalUnbondingDelegation(cdc, kvA.Value)
+			undelegationB := types.MustUnmarshalUnbondingDelegation(cdc, kvB.Value)
+			return fmt.Sprintf("%v\n%v", undelegationA, undelegationB)
+
+		case bytes.Equal(kvA.Key[:1], types.ServiceDelegationPrefix):
+			delegationA := types.MustUnmarshalDelegation(cdc, kvA.Value)
+			delegationB := types.MustUnmarshalDelegation(cdc, kvB.Value)
+			return fmt.Sprintf("%v\n%v", delegationA, delegationB)
+
+		case bytes.Equal(kvA.Key[:1], types.ServiceUnbondingDelegationPrefix):
+			undelegationA := types.MustUnmarshalUnbondingDelegation(cdc, kvA.Value)
+			undelegationB := types.MustUnmarshalUnbondingDelegation(cdc, kvB.Value)
+			return fmt.Sprintf("%v\n%v", undelegationA, undelegationB)
+
+		case bytes.Equal(kvA.Key[:1], types.UnbondingQueueKey):
+			var listA, listB types.DTDataList
+			cdc.MustUnmarshal(kvA.Value, &listA)
+			cdc.MustUnmarshal(kvB.Value, &listB)
+			return fmt.Sprintf("%v\n%v", listA, listB)
+
+		// Collections
+		case bytes.Equal(kvA.Key[:1], types.OperatorJoinedServicesPrefix),
+			bytes.Equal(kvA.Key[:1], types.ServiceOperatorsAllowListPrefix),
+			bytes.Equal(kvA.Key[:1], types.ServiceSecuringPoolsPrefix),
+			bytes.Equal(kvA.Key[:1], types.UserPreferencesPrefix):
+			return collectionsDecoder(kvA, kvB)
+
+		default:
+			panic(fmt.Sprintf("invalid restaking key prefix %X", kvA.Key[:1]))
+		}
+	}
+}

x/restaking/simulation/genesis.go

@@ -0,0 +1,30 @@
+package simulation
+
+import (
+	"github.com/cosmos/cosmos-sdk/types/module"
+
+	operatorssimulation "github.com/milkyway-labs/milkyway/v3/x/operators/simulation"
+	poolssimulation "github.com/milkyway-labs/milkyway/v3/x/pools/simulation"
+	"github.com/milkyway-labs/milkyway/v3/x/restaking/types"
+	servicessimulation "github.com/milkyway-labs/milkyway/v3/x/services/simulation"
+)
+
+// RandomizedGenState generates a random GenesisState for the restaking module
+func RandomizedGenState(simState *module.SimulationState) {
+	poolsGenesis := poolssimulation.GetGenesisState(simState)
+	operatorsGenesis := operatorssimulation.GetGenesisState(simState)
+	servicesGenesis := servicessimulation.GetGenesisState(simState)
+
+	genesis := types.NewGenesis(
+		RandomOperatorJoinedServices(simState.Rand, operatorsGenesis.Operators, servicesGenesis.Services),
+		RandomServiceAllowedOperators(simState.Rand, servicesGenesis.Services, operatorsGenesis.Operators),
+		RandomServiceSecuringPools(simState.Rand, poolsGenesis.Pools, servicesGenesis.Services),
+		// empty delegations and undelegations since we need to also perform side effects to other
+		// modules to keep the shares consistent.
+		nil,
+		nil,
+		RandomUserPreferencesEntries(simState.Rand, servicesGenesis.Services),
+		RandomParams(simState.Rand),
+	)
+	simState.GenState[types.ModuleName] = simState.Cdc.MustMarshalJSON(genesis)
+}