fix: simulation tests

app/modules.go

@@ -154,7 +154,7 @@ func appModules(
 		pools.NewAppModule(appCodec, app.PoolsKeeper),
 		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, app.AccountKeeper, app.BankKeeper, app.OperatorsKeeper, app.ServicesKeeper),
+		rewards.NewAppModule(appCodec, app.RewardsKeeper, app.AccountKeeper, app.BankKeeper, app.PoolsKeeper, app.OperatorsKeeper, app.ServicesKeeper),
 		liquidvesting.NewAppModule(appCodec, app.LiquidVestingKeeper),
 	}
 }
@@ -206,7 +206,7 @@ func simulationModules(
 		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),
-		rewards.NewAppModule(appCodec, app.RewardsKeeper, app.AccountKeeper, app.BankKeeper, app.OperatorsKeeper, app.ServicesKeeper),
+		rewards.NewAppModule(appCodec, app.RewardsKeeper, app.AccountKeeper, app.BankKeeper, app.PoolsKeeper, app.OperatorsKeeper, app.ServicesKeeper),
 	}
 }
 

testutils/simtesting/utils.go

@@ -88,6 +88,17 @@ func RandomCoin(r *rand.Rand, denom string, maxAmount int) sdk.Coin {
 	)
 }
 
+// RandomDecCoins returns a random list of DecCoins by randomly selecting
+// a sub set of the provided denoms.
+func RandomDecCoins(r *rand.Rand, denoms []string, maxAmount sdkmath.LegacyDec) sdk.DecCoins {
+	coins := sdk.NewDecCoins()
+	for _, denom := range RandomSubSlice(r, denoms) {
+		coins = coins.Add(sdk.NewDecCoinFromDec(denom, simtypes.RandomDecAmount(r, maxAmount)))
+	}
+
+	return coins
+}
+
 // 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
@@ -135,3 +146,8 @@ func RandomPositiveUint64(r *rand.Rand) uint64 {
 	}
 	return value
 }
+
+// RandomSliceElement returns a random element from the given slice
+func RandomSliceElement[T any](r *rand.Rand, slice []T) T {
+	return slice[r.Intn(len(slice))]
+}

x/rewards/module.go

@@ -18,6 +18,7 @@ import (
 	"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/rewards/client/cli"
 	"github.com/milkyway-labs/milkyway/v3/x/rewards/keeper"
 	"github.com/milkyway-labs/milkyway/v3/x/rewards/simulation"
@@ -103,6 +104,7 @@ type AppModule struct {
 	accountKeeper authkeeper.AccountKeeper
 	bankKeeper    bankkeeper.Keeper
 
+	poolsKeeper     *poolskeeper.Keeper
 	operatorsKeeper *operatorskeeper.Keeper
 	servicesKeeper  *serviceskeeper.Keeper
 }
@@ -112,6 +114,7 @@ func NewAppModule(
 	keeper *keeper.Keeper,
 	accountKeeper authkeeper.AccountKeeper,
 	bankKeeper bankkeeper.Keeper,
+	poolsKeeper *poolskeeper.Keeper,
 	operatorsKeeper *operatorskeeper.Keeper,
 	serviceKeeper *serviceskeeper.Keeper,
 ) AppModule {
@@ -120,6 +123,7 @@ func NewAppModule(
 		keeper:          keeper,
 		accountKeeper:   accountKeeper,
 		bankKeeper:      bankKeeper,
+		poolsKeeper:     poolsKeeper,
 		operatorsKeeper: operatorsKeeper,
 		servicesKeeper:  serviceKeeper,
 	}
@@ -198,6 +202,7 @@ func (am AppModule) WeightedOperations(simState module.SimulationState) []simtyp
 		simState.AppParams,
 		am.accountKeeper,
 		am.bankKeeper,
+		am.poolsKeeper,
 		am.operatorsKeeper,
 		am.servicesKeeper,
 		am.keeper,

x/rewards/simulation/genesis.go

@@ -3,24 +3,25 @@ package simulation
 import (
 	"github.com/cosmos/cosmos-sdk/types/module"
 
-	poolstypes "github.com/milkyway-labs/milkyway/v3/x/pools/types"
+	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/rewards/types"
-	servicestypes "github.com/milkyway-labs/milkyway/v3/x/services/types"
+	servicessimulation "github.com/milkyway-labs/milkyway/v3/x/services/simulation"
 )
 
 // RandomizedGenState generates a random GenesisState for the operators module
 func RandomizedGenState(simState *module.SimulationState) {
-	// Get the services genesis state
-	rawServicesGenesis := simState.GenState[servicestypes.ModuleName]
-	var servicesGenesis servicestypes.GenesisState
-	simState.Cdc.MustUnmarshalJSON(rawServicesGenesis, &servicesGenesis)
+	servicesGenesis := servicessimulation.GetGenesisState(simState)
+	poolsGenesis := poolssimulation.GetGenesisState(simState)
+	operatorsGenesis := operatorssimulation.GetGenesisState(simState)
 
-	// Get the pools genesis state
-	rawPoolsGenesis := simState.GenState[poolstypes.ModuleName]
-	var poolsGenesis poolstypes.GenesisState
-	simState.Cdc.MustUnmarshalJSON(rawPoolsGenesis, &poolsGenesis)
-
-	rewardsPlans := RandomRewardsPlans(simState.Rand, servicesGenesis.Services, []string{simState.BondDenom})
+	rewardsPlans := RandomRewardsPlans(
+		simState.Rand,
+		poolsGenesis.Pools,
+		operatorsGenesis.Operators,
+		servicesGenesis.Services,
+		[]string{simState.BondDenom},
+	)
 	nextRewardsPlan := uint64(1)
 	for _, plan := range rewardsPlans {
 		if plan.ID >= nextRewardsPlan {
@@ -39,7 +40,7 @@ func RandomizedGenState(simState *module.SimulationState) {
 		types.NewDelegationTypeRecords(nil, nil, nil, nil),
 		types.NewDelegationTypeRecords(nil, nil, nil, nil),
 		types.NewDelegationTypeRecords(nil, nil, nil, nil),
-		RandomOperatorAccumulatedCommissionRecords(simState.Rand, []string{simState.BondDenom}),
+		RandomOperatorAccumulatedCommissionRecords(simState.Rand, operatorsGenesis.Operators, []string{simState.BondDenom}),
 		RandomPoolServiceTotalDelegatorShares(simState.Rand, poolsGenesis, servicesGenesis, []string{simState.BondDenom}),
 	)
 	simState.GenState[types.ModuleName] = simState.Cdc.MustMarshalJSON(genesis)

x/rewards/simulation/msg_factory.go

@@ -16,6 +16,9 @@ import (
 	"github.com/milkyway-labs/milkyway/v3/utils"
 	operatorskeeper "github.com/milkyway-labs/milkyway/v3/x/operators/keeper"
 	operatorssimulation "github.com/milkyway-labs/milkyway/v3/x/operators/simulation"
+	operatorstypes "github.com/milkyway-labs/milkyway/v3/x/operators/types"
+	poolskeeper "github.com/milkyway-labs/milkyway/v3/x/pools/keeper"
+	poolssimulation "github.com/milkyway-labs/milkyway/v3/x/pools/simulation"
 	restakingtypes "github.com/milkyway-labs/milkyway/v3/x/restaking/types"
 	"github.com/milkyway-labs/milkyway/v3/x/rewards/keeper"
 	"github.com/milkyway-labs/milkyway/v3/x/rewards/types"
@@ -42,6 +45,7 @@ func WeightedOperations(
 	appParams simtypes.AppParams,
 	ak authkeeper.AccountKeeper,
 	bk bankkeeper.Keeper,
+	pk *poolskeeper.Keeper,
 	ok *operatorskeeper.Keeper,
 	sk *serviceskeeper.Keeper,
 	k *keeper.Keeper,
@@ -76,8 +80,8 @@ func WeightedOperations(
 	})
 
 	return simulation.WeightedOperations{
-		simulation.NewWeightedOperation(weightMsgCreateRewardsPlan, SimulateMsgCreateRewardsPlan(ak, bk, sk, k)),
-		simulation.NewWeightedOperation(weightMsgEditRewardsPlan, SimulateMsgEditRewardsPlan(ak, bk, sk, k)),
+		simulation.NewWeightedOperation(weightMsgCreateRewardsPlan, SimulateMsgCreateRewardsPlan(ak, bk, pk, ok, sk, k)),
+		simulation.NewWeightedOperation(weightMsgEditRewardsPlan, SimulateMsgEditRewardsPlan(ak, bk, pk, ok, sk, k)),
 		simulation.NewWeightedOperation(weightMsgSetWithdrawAddress, SimulateMsgSetWithdrawAddress(ak, bk, k)),
 		simulation.NewWeightedOperation(weightMsgWithdrawDelegatorReward, SimulateMsgWithdrawDelegatorReward(ak, bk, k)),
 		simulation.NewWeightedOperation(weightMsgWithdrawOperatorCommission, SimulateMsgWithdrawOperatorCommission(ak, bk, ok, k)),
@@ -89,6 +93,8 @@ func WeightedOperations(
 func SimulateMsgCreateRewardsPlan(
 	ak authkeeper.AccountKeeper,
 	bk bankkeeper.Keeper,
+	pk *poolskeeper.Keeper,
+	ok *operatorskeeper.Keeper,
 	sk *serviceskeeper.Keeper,
 	k *keeper.Keeper,
 ) simtypes.Operation {
@@ -129,6 +135,20 @@ func SimulateMsgCreateRewardsPlan(
 			), nil, nil
 		}
 
+		// Get a random pool that we will use to configure the pool distribution
+		pool, found := poolssimulation.GetRandomExistingPool(r, ctx, pk, nil)
+		if !found {
+			return simtypes.NoOpMsg(types.ModuleName, sdk.MsgTypeURL(msg), "no pool found"), nil, nil
+		}
+
+		// Get a random operator that we will use to configure the operator distribution
+		operator, found := operatorssimulation.GetRandomExistingOperator(r, ctx, ok, func(o operatorstypes.Operator) bool {
+			return o.IsActive()
+		})
+		if !found {
+			return simtypes.NoOpMsg(types.ModuleName, sdk.MsgTypeURL(msg), "no operator found"), nil, nil
+		}
+
 		// Compute some random start/end time
 		rewardsStart := ctx.BlockTime().Add(time.Hour * time.Duration(r.Intn(10)+1))
 		rewardsEnd := rewardsStart.Add(time.Hour * time.Duration(r.Intn(96)+1))
@@ -150,8 +170,8 @@ func SimulateMsgCreateRewardsPlan(
 			amount,
 			rewardsStart,
 			rewardsEnd,
-			RandomDistribution(r, restakingtypes.DELEGATION_TYPE_POOL),
-			RandomDistribution(r, restakingtypes.DELEGATION_TYPE_OPERATOR),
+			RandomDistribution(r, restakingtypes.DELEGATION_TYPE_POOL, pool),
+			RandomDistribution(r, restakingtypes.DELEGATION_TYPE_OPERATOR, operator),
 			RandomUsersDistribution(r),
 			rewardsParams.RewardsPlanCreationFee,
 			service.Admin,
@@ -163,6 +183,8 @@ func SimulateMsgCreateRewardsPlan(
 func SimulateMsgEditRewardsPlan(
 	ak authkeeper.AccountKeeper,
 	bk bankkeeper.Keeper,
+	pk *poolskeeper.Keeper,
+	ok *operatorskeeper.Keeper,
 	sk *serviceskeeper.Keeper,
 	k *keeper.Keeper,
 ) simtypes.Operation {
@@ -176,6 +198,20 @@ func SimulateMsgEditRewardsPlan(
 			return simtypes.NoOpMsg(types.ModuleName, sdk.MsgTypeURL(msg), "no rewards plan found"), nil, nil
 		}
 
+		// Get a random pool that we will use to configure the pool distribution
+		pool, found := poolssimulation.GetRandomExistingPool(r, ctx, pk, nil)
+		if !found {
+			return simtypes.NoOpMsg(types.ModuleName, sdk.MsgTypeURL(msg), "no pool found"), nil, nil
+		}
+
+		// Get a random operator that we will use to configure the operator distribution
+		operator, found := operatorssimulation.GetRandomExistingOperator(r, ctx, ok, func(o operatorstypes.Operator) bool {
+			return o.IsActive()
+		})
+		if !found {
+			return simtypes.NoOpMsg(types.ModuleName, sdk.MsgTypeURL(msg), "no operator found"), nil, nil
+		}
+
 		// Get the service admin
 		service, found, err := sk.GetService(ctx, plan.ServiceID)
 		if err != nil {
@@ -212,8 +248,8 @@ func SimulateMsgEditRewardsPlan(
 			amount,
 			rewardsStart,
 			rewardsEnd,
-			RandomDistribution(r, restakingtypes.DELEGATION_TYPE_POOL),
-			RandomDistribution(r, restakingtypes.DELEGATION_TYPE_OPERATOR),
+			RandomDistribution(r, restakingtypes.DELEGATION_TYPE_POOL, pool),
+			RandomDistribution(r, restakingtypes.DELEGATION_TYPE_OPERATOR, operator),
 			RandomUsersDistribution(r),
 			service.Admin,
 		)