test: add x/services simulation tests

app/keepers/keepers.go

@@ -275,8 +275,6 @@ func NewAppKeeper(
 		logger,
 	)
 
-	communityPoolKeeper := NewCommunityPoolKeeper(appKeepers.BankKeeper, authtypes.FeeCollectorName)
-
 	appKeepers.CrisisKeeper = crisiskeeper.NewKeeper(
 		appCodec,
 		runtime.NewKVStoreService(appKeepers.keys[crisistypes.StoreKey]),
@@ -421,7 +419,7 @@ func NewAppKeeper(
 		runtime.NewKVStoreService(appKeepers.keys[tokenfactorytypes.StoreKey]),
 		appKeepers.AccountKeeper,
 		appKeepers.BankKeeper,
-		communityPoolKeeper,
+		appKeepers.DistrKeeper,
 		authtypes.NewModuleAddress(govtypes.ModuleName).String(),
 	)
 	appKeepers.TokenFactoryKeeper.SetContractKeeper(contractKeeper)
@@ -580,15 +578,15 @@ func NewAppKeeper(
 		appKeepers.keys[servicestypes.StoreKey],
 		runtime.NewKVStoreService(appKeepers.keys[servicestypes.StoreKey]),
 		appKeepers.AccountKeeper,
-		communityPoolKeeper,
+		appKeepers.DistrKeeper,
 		govAuthority,
 	)
 	appKeepers.OperatorsKeeper = operatorskeeper.NewKeeper(
 		appCodec,
 		appKeepers.keys[operatorstypes.StoreKey],
 		runtime.NewKVStoreService(appKeepers.keys[operatorstypes.StoreKey]),
 		appKeepers.AccountKeeper,
-		communityPoolKeeper,
+		appKeepers.DistrKeeper,
 		govAuthority,
 	)
 	appKeepers.PoolsKeeper = poolskeeper.NewKeeper(
@@ -608,11 +606,6 @@ func NewAppKeeper(
 		appKeepers.ServicesKeeper,
 		govAuthority,
 	)
-
-	// Set hooks based on the restaking keeper
-	appKeepers.OperatorsKeeper.SetHooks(appKeepers.RestakingKeeper.OperatorsHooks())
-	appKeepers.ServicesKeeper.SetHooks(appKeepers.RestakingKeeper.ServicesHooks())
-
 	appKeepers.AssetsKeeper = assetskeeper.NewKeeper(
 		appCodec,
 		runtime.NewKVStoreService(appKeepers.keys[assetstypes.StoreKey]),
@@ -623,7 +616,7 @@ func NewAppKeeper(
 		runtime.NewKVStoreService(appKeepers.keys[rewardstypes.StoreKey]),
 		appKeepers.AccountKeeper,
 		appKeepers.BankKeeper,
-		communityPoolKeeper,
+		appKeepers.DistrKeeper,
 		appKeepers.OracleKeeper,
 		appKeepers.PoolsKeeper,
 		appKeepers.OperatorsKeeper,
@@ -634,8 +627,6 @@ func NewAppKeeper(
 	)
 
 	// Set hooks based on the rewards keeper
-	appKeepers.RestakingKeeper.SetHooks(appKeepers.RewardsKeeper.Hooks())
-
 	appKeepers.LiquidVestingKeeper = liquidvestingkeeper.NewKeeper(
 		appCodec,
 		appKeepers.keys[liquidvestingtypes.StoreKey],
@@ -650,9 +641,25 @@ func NewAppKeeper(
 		govAuthority,
 	)
 
-	// Set hooks based on the liquid vesting keeper
+	// Set the restrictions on sending tokens
 	appKeepers.BankKeeper.AppendSendRestriction(appKeepers.LiquidVestingKeeper.SendRestrictionFn)
 
+	// Set the hooks up to this point
+	appKeepers.PoolsKeeper.SetHooks(
+		appKeepers.RewardsKeeper.PoolsHooks(),
+	)
+	appKeepers.OperatorsKeeper.SetHooks(operatorstypes.NewMultiOperatorsHooks(
+		appKeepers.RestakingKeeper.OperatorsHooks(),
+		appKeepers.RewardsKeeper.OperatorsHooks(),
+	))
+	appKeepers.ServicesKeeper.SetHooks(servicestypes.NewMultiServicesHooks(
+		appKeepers.RestakingKeeper.ServicesHooks(),
+		appKeepers.RewardsKeeper.ServicesHooks(),
+	))
+	appKeepers.RestakingKeeper.SetHooks(
+		appKeepers.RewardsKeeper.RestakingHooks(),
+	)
+
 	// ---------------------- //
 	// --- Stride Keepers --- //
 	// ---------------------- //

app/modules.go

@@ -176,7 +176,7 @@ func appModules(
 		icacallbacks.NewAppModule(appCodec, *app.ICACallbacksKeeper, app.AccountKeeper, app.BankKeeper),
 
 		// MilkyWay modules
-		services.NewAppModule(appCodec, app.ServicesKeeper, app.PoolsKeeper),
+		services.NewAppModule(appCodec, app.ServicesKeeper, app.PoolsKeeper, app.AccountKeeper, app.BankKeeper),
 		operators.NewAppModule(appCodec, app.OperatorsKeeper),
 		pools.NewAppModule(appCodec, app.PoolsKeeper),
 		restaking.NewAppModule(appCodec, app.RestakingKeeper),
@@ -229,6 +229,9 @@ func simulationModules(
 		ibc.NewAppModule(app.IBCKeeper),
 		app.TransferModule,
 		app.ICAModule,
+
+		// MilkyWay modules
+		services.NewAppModule(appCodec, app.ServicesKeeper, app.PoolsKeeper, app.AccountKeeper, app.BankKeeper),
 	}
 }
 

testutils/simtesting/utils.go

@@ -0,0 +1,63 @@
+package simtesting
+
+import (
+	"math/rand"
+
+	"github.com/cosmos/cosmos-sdk/baseapp"
+	"github.com/cosmos/cosmos-sdk/codec"
+	codectypes "github.com/cosmos/cosmos-sdk/codec/types"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	simtypes "github.com/cosmos/cosmos-sdk/types/simulation"
+	authkeeper "github.com/cosmos/cosmos-sdk/x/auth/keeper"
+	"github.com/cosmos/cosmos-sdk/x/auth/tx"
+	bankkeeper "github.com/cosmos/cosmos-sdk/x/bank/keeper"
+	"github.com/cosmos/cosmos-sdk/x/simulation"
+)
+
+// SendMsg sends a transaction with the specified message.
+func SendMsg(
+	r *rand.Rand, moduleName string, app *baseapp.BaseApp, ak authkeeper.AccountKeeper, bk bankkeeper.Keeper,
+	msg sdk.Msg, ctx sdk.Context,
+	simAccount simtypes.Account,
+) (simtypes.OperationMsg, []simtypes.FutureOperation, error) {
+	deposit := sdk.Coins{}
+	spendableCoins := bk.SpendableCoins(ctx, simAccount.Address)
+	for _, v := range spendableCoins {
+		if bk.IsSendEnabledCoin(ctx, v) {
+			deposit = deposit.Add(simtypes.RandSubsetCoins(r, sdk.NewCoins(v))...)
+		}
+	}
+
+	if deposit.IsZero() {
+		msgType := sdk.MsgTypeURL(msg)
+		return simtypes.NoOpMsg(moduleName, msgType, "skip because of broke account"), nil, nil
+	}
+
+	interfaceRegistry := codectypes.NewInterfaceRegistry()
+	txConfig := tx.NewTxConfig(codec.NewProtoCodec(interfaceRegistry), tx.DefaultSignModes)
+	txCtx := simulation.OperationInput{
+		R:               r,
+		App:             app,
+		TxGen:           txConfig,
+		Cdc:             nil,
+		Msg:             msg,
+		Context:         ctx,
+		SimAccount:      simAccount,
+		AccountKeeper:   ak,
+		Bankkeeper:      bk,
+		ModuleName:      moduleName,
+		CoinsSpentInMsg: deposit,
+	}
+	return simulation.GenAndDeliverTxWithRandFees(txCtx)
+}
+
+// GetSimAccount gets the Account with the given address
+func GetSimAccount(address sdk.Address, accs []simtypes.Account) (simtypes.Account, bool) {
+	for _, acc := range accs {
+		if acc.Address.Equals(address) {
+			return acc, true
+		}
+	}
+	return simtypes.Account{}, false
+}

testutils/storetesting/keeper.go

@@ -13,6 +13,8 @@ import (
 	authkeeper "github.com/cosmos/cosmos-sdk/x/auth/keeper"
 	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
 	banktypes "github.com/cosmos/cosmos-sdk/x/bank/types"
+	distrkeeper "github.com/cosmos/cosmos-sdk/x/distribution/keeper"
+	distrtypes "github.com/cosmos/cosmos-sdk/x/distribution/types"
 	govtypes "github.com/cosmos/cosmos-sdk/x/gov/types"
 
 	milkyway "github.com/milkyway-labs/milkyway/app"
@@ -28,8 +30,9 @@ type BaseKeeperTestData struct {
 
 	AuthorityAddress string
 
-	AccountKeeper authkeeper.AccountKeeper
-	BankKeeper    bankkeeper.Keeper
+	AccountKeeper      authkeeper.AccountKeeper
+	BankKeeper         bankkeeper.Keeper
+	DistributionKeeper distrkeeper.Keeper
 }
 
 // NewBaseKeeperTestData returns a new BaseKeeperTestData
@@ -45,7 +48,7 @@ func NewBaseKeeperTestData(t *testing.T, keys []string) BaseKeeperTestData {
 	var data BaseKeeperTestData
 
 	// Define store keys
-	keys = append(keys, []string{authtypes.StoreKey, banktypes.StoreKey}...)
+	keys = append(keys, []string{authtypes.StoreKey, banktypes.StoreKey, distrtypes.StoreKey}...)
 	slices.Sort(keys)
 	keys = slices.Compact(keys)
 	data.Keys = storetypes.NewKVStoreKeys(keys...)
@@ -77,6 +80,20 @@ func NewBaseKeeperTestData(t *testing.T, keys []string) BaseKeeperTestData {
 		data.AuthorityAddress,
 		log.NewNopLogger(),
 	)
+	data.DistributionKeeper = distrkeeper.NewKeeper(
+		data.Cdc,
+		runtime.NewKVStoreService(data.Keys[distrtypes.StoreKey]),
+		data.AccountKeeper,
+		data.BankKeeper,
+		nil,
+		authtypes.FeeCollectorName,
+		data.AuthorityAddress,
+	)
+
+	// Init the module's genesis state as the default ones
+	data.AccountKeeper.InitGenesis(data.Context, *authtypes.DefaultGenesisState())
+	data.BankKeeper.InitGenesis(data.Context, banktypes.DefaultGenesisState())
+	data.DistributionKeeper.InitGenesis(data.Context, *distrtypes.DefaultGenesisState())
 
 	return data
 }

utils/slices.go

@@ -80,3 +80,14 @@ func Intersect[T comparable](a, b []T) []T {
 
 	return result
 }
+
+// Filter returns the elements of the slice that satisfy the given predicate.
+func Filter[T any](slice []T, f func(T) bool) []T {
+	var result []T
+	for _, v := range slice {
+		if f(v) {
+			result = append(result, v)
+		}
+	}
+	return result
+}

x/liquidvesting/keeper/common_test.go

@@ -143,7 +143,7 @@ func (suite *KeeperTestSuite) createService(ctx sdk.Context, id uint32) {
 }
 
 func (suite *KeeperTestSuite) createOperator(ctx sdk.Context, id uint32) {
-	err := suite.ok.RegisterOperator(ctx, operatorstypes.NewOperator(
+	err := suite.ok.CreateOperator(ctx, operatorstypes.NewOperator(
 		id,
 		operatorstypes.OPERATOR_STATUS_ACTIVE,
 		fmt.Sprintf("operator-%d", id),

x/liquidvesting/testutils/keeper.go

@@ -15,7 +15,6 @@ import (
 
 	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
 
-	appkeepers "github.com/milkyway-labs/milkyway/app/keepers"
 	"github.com/milkyway-labs/milkyway/testutils/storetesting"
 	"github.com/milkyway-labs/milkyway/x/liquidvesting"
 	"github.com/milkyway-labs/milkyway/x/liquidvesting/keeper"
@@ -68,24 +67,20 @@ func NewKeeperTestData(t *testing.T) KeeperTestData {
 		runtime.NewKVStoreService(data.Keys[poolstypes.StoreKey]),
 		data.AccountKeeper,
 	)
-	communityPoolKeeper := appkeepers.NewCommunityPoolKeeper(
-		data.BankKeeper,
-		authtypes.FeeCollectorName,
-	)
 	data.OperatorsKeeper = operatorskeeper.NewKeeper(
 		data.Cdc,
 		data.Keys[operatorstypes.StoreKey],
 		runtime.NewKVStoreService(data.Keys[operatorstypes.StoreKey]),
 		data.AccountKeeper,
-		communityPoolKeeper,
+		data.DistributionKeeper,
 		data.AuthorityAddress,
 	)
 	data.ServicesKeeper = serviceskeeper.NewKeeper(
 		data.Cdc,
 		data.Keys[servicestypes.StoreKey],
 		runtime.NewKVStoreService(data.Keys[servicestypes.StoreKey]),
 		data.AccountKeeper,
-		communityPoolKeeper,
+		data.DistributionKeeper,
 		data.AuthorityAddress,
 	)
 	data.RestakingKeeper = restakingkeeper.NewKeeper(

x/operators/keeper/genesis.go

@@ -47,7 +47,7 @@ func (k *Keeper) InitGenesis(ctx sdk.Context, state types.GenesisState) error {
 
 	// Store the operators
 	for _, operator := range state.Operators {
-		if err := k.SaveOperator(ctx, operator); err != nil {
+		if err := k.CreateOperator(ctx, operator); err != nil {
 			return err
 		}
 	}

x/operators/keeper/genesis_test.go

@@ -35,7 +35,7 @@ func (suite *KeeperTestSuite) TestKeeper_ExportGenesis() {
 				suite.k.SetNextOperatorID(ctx, 10)
 				suite.k.SetParams(ctx, types.DefaultParams())
 
-				err := suite.k.RegisterOperator(ctx, types.NewOperator(
+				err := suite.k.CreateOperator(ctx, types.NewOperator(
 					1,
 					types.OPERATOR_STATUS_ACTIVE,
 					"MilkyWay Operator",
@@ -45,7 +45,7 @@ func (suite *KeeperTestSuite) TestKeeper_ExportGenesis() {
 				))
 				suite.Require().NoError(err)
 
-				err = suite.k.RegisterOperator(ctx, types.NewOperator(
+				err = suite.k.CreateOperator(ctx, types.NewOperator(
 					2,
 					types.OPERATOR_STATUS_INACTIVATING,
 					"Inertia",
@@ -95,13 +95,13 @@ func (suite *KeeperTestSuite) TestKeeper_ExportGenesis() {
 					"https://milkyway.com/picture",
 					"cosmos167x6ehhple8gwz5ezy9x0464jltvdpzl6qfdt4",
 				)
-				err := suite.k.RegisterOperator(ctx, activeValidator)
+				err := suite.k.CreateOperator(ctx, activeValidator)
 				suite.Require().NoError(err)
 
 				err = suite.k.StartOperatorInactivation(ctx, activeValidator)
 				suite.Require().NoError(err)
 
-				err = suite.k.RegisterOperator(ctx, types.NewOperator(
+				err = suite.k.CreateOperator(ctx, types.NewOperator(
 					2,
 					types.OPERATOR_STATUS_ACTIVE,
 					"Inertia",
@@ -146,7 +146,7 @@ func (suite *KeeperTestSuite) TestKeeper_ExportGenesis() {
 				suite.k.SetNextOperatorID(ctx, 10)
 				suite.k.SetParams(ctx, types.DefaultParams())
 
-				err := suite.k.RegisterOperator(ctx, types.NewOperator(
+				err := suite.k.CreateOperator(ctx, types.NewOperator(
 					1,
 					types.OPERATOR_STATUS_ACTIVE,
 					"MilkyWay Operator",
@@ -162,7 +162,7 @@ func (suite *KeeperTestSuite) TestKeeper_ExportGenesis() {
 				))
 				suite.Require().NoError(err)
 
-				err = suite.k.RegisterOperator(ctx, types.NewOperator(
+				err = suite.k.CreateOperator(ctx, types.NewOperator(
 					2,
 					types.OPERATOR_STATUS_INACTIVATING,
 					"Inertia",