fix: restaking simulation tests

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))
+	for i := 0; i < count; i++ {
+		// 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/module.go

@@ -18,10 +18,13 @@ import (
 	"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 (
@@ -105,19 +108,19 @@ type AppModule struct {
 
 	ak              authkeeper.AccountKeeper
 	bk              bankkeeper.Keeper
-	poolsKeeper     types.PoolsKeeper
-	operatorsKeeper types.OperatorsKeeper
-	servicesKeeper  types.ServicesKeeper
+	poolsKeeper     *poolskeeper.Keeper
+	operatorsKeeper *operatorskeeper.Keeper
+	servicesKeeper  *serviceskeeper.Keeper
 }
 
 func NewAppModule(
 	cdc codec.Codec,
 	keeper *keeper.Keeper,
 	ak authkeeper.AccountKeeper,
 	bk bankkeeper.Keeper,
-	poolsKeeper types.PoolsKeeper,
-	operatorsKeeper types.OperatorsKeeper,
-	servicesKeeper types.ServicesKeeper,
+	poolsKeeper *poolskeeper.Keeper,
+	operatorsKeeper *operatorskeeper.Keeper,
+	servicesKeeper *serviceskeeper.Keeper,
 ) AppModule {
 	return AppModule{
 		AppModuleBasic:  NewAppModuleBasic(cdc),

x/restaking/simulation/genesis.go

@@ -3,21 +3,28 @@ 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),
-		RandomServiceAllowedOperators(simState),
-		RandomServiceSecuringPools(simState),
+		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),
-		RandomParams(simState),
+		RandomUserPreferencesEntries(simState.Rand, servicesGenesis.Services),
+		RandomParams(simState.Rand),
 	)
 	simState.GenState[types.ModuleName] = simState.Cdc.MustMarshalJSON(genesis)
 }