feat(x/protocolpool): add env bundler to protocolpool module (#19420)

Co-authored-by: Marko <marbar3778@yahoo.com>

simapp/app.go

@@ -337,7 +337,7 @@ func NewSimApp(
 	)
 	app.MintKeeper = mintkeeper.NewKeeper(appCodec, runtime.NewEnvironment(runtime.NewKVStoreService(keys[minttypes.StoreKey]), logger), app.StakingKeeper, app.AuthKeeper, app.BankKeeper, authtypes.FeeCollectorName, authtypes.NewModuleAddress(govtypes.ModuleName).String())
 
-	app.PoolKeeper = poolkeeper.NewKeeper(appCodec, runtime.NewKVStoreService(keys[pooltypes.StoreKey]), app.AuthKeeper, app.BankKeeper, app.StakingKeeper, authtypes.NewModuleAddress(govtypes.ModuleName).String())
+	app.PoolKeeper = poolkeeper.NewKeeper(appCodec, runtime.NewEnvironment(runtime.NewKVStoreService(keys[pooltypes.StoreKey]), logger), app.AuthKeeper, app.BankKeeper, app.StakingKeeper, authtypes.NewModuleAddress(govtypes.ModuleName).String())
 
 	app.DistrKeeper = distrkeeper.NewKeeper(appCodec, runtime.NewKVStoreService(keys[distrtypes.StoreKey]), app.AuthKeeper, app.BankKeeper, app.StakingKeeper, app.PoolKeeper, authtypes.FeeCollectorName, authtypes.NewModuleAddress(govtypes.ModuleName).String())
 

tests/integration/distribution/keeper/msg_server_test.go

@@ -108,9 +108,7 @@ func initFixture(t *testing.T) *fixture {
 	stakingKeeper := stakingkeeper.NewKeeper(cdc, runtime.NewEnvironment(runtime.NewKVStoreService(keys[stakingtypes.StoreKey]), log.NewNopLogger()), accountKeeper, bankKeeper, authority.String(), addresscodec.NewBech32Codec(sdk.Bech32PrefixValAddr), addresscodec.NewBech32Codec(sdk.Bech32PrefixConsAddr))
 	require.NoError(t, stakingKeeper.Params.Set(newCtx, stakingtypes.DefaultParams()))
 
-	poolKeeper := poolkeeper.NewKeeper(
-		cdc, runtime.NewKVStoreService(keys[pooltypes.StoreKey]), accountKeeper, bankKeeper, stakingKeeper, authority.String(),
-	)
+	poolKeeper := poolkeeper.NewKeeper(cdc, runtime.NewEnvironment(runtime.NewKVStoreService(keys[pooltypes.StoreKey]), log.NewNopLogger()), accountKeeper, bankKeeper, stakingKeeper, authority.String())
 
 	distrKeeper := distrkeeper.NewKeeper(
 		cdc, runtime.NewKVStoreService(keys[distrtypes.StoreKey]), accountKeeper, bankKeeper, stakingKeeper, poolKeeper, distrtypes.ModuleName, authority.String(),

tests/integration/gov/keeper/keeper_test.go

@@ -94,7 +94,7 @@ func initFixture(tb testing.TB) *fixture {
 
 	stakingKeeper := stakingkeeper.NewKeeper(cdc, runtime.NewEnvironment(runtime.NewKVStoreService(keys[stakingtypes.StoreKey]), log.NewNopLogger()), accountKeeper, bankKeeper, authority.String(), addresscodec.NewBech32Codec(sdk.Bech32PrefixValAddr), addresscodec.NewBech32Codec(sdk.Bech32PrefixConsAddr))
 
-	poolKeeper := poolkeeper.NewKeeper(cdc, runtime.NewKVStoreService(keys[pooltypes.StoreKey]), accountKeeper, bankKeeper, stakingKeeper, authority.String())
+	poolKeeper := poolkeeper.NewKeeper(cdc, runtime.NewEnvironment(runtime.NewKVStoreService(keys[pooltypes.StoreKey]), log.NewNopLogger()), accountKeeper, bankKeeper, stakingKeeper, authority.String())
 
 	// set default staking params
 	err := stakingKeeper.Params.Set(newCtx, stakingtypes.DefaultParams())

x/protocolpool/depinject.go

@@ -3,7 +3,6 @@ package protocolpool
 import (
 	modulev1 "cosmossdk.io/api/cosmos/protocolpool/module/v1"
 	"cosmossdk.io/core/appmodule"
-	storetypes "cosmossdk.io/core/store"
 	"cosmossdk.io/depinject"
 	"cosmossdk.io/depinject/appconfig"
 	authtypes "cosmossdk.io/x/auth/types"
@@ -31,9 +30,9 @@ func init() {
 type ModuleInputs struct {
 	depinject.In
 
-	Config       *modulev1.Module
-	Codec        codec.Codec
-	StoreService storetypes.KVStoreService
+	Config      *modulev1.Module
+	Codec       codec.Codec
+	Environment appmodule.Environment
 
 	AccountKeeper types.AccountKeeper
 	BankKeeper    types.BankKeeper
@@ -54,7 +53,7 @@ func ProvideModule(in ModuleInputs) ModuleOutputs {
 		authority = authtypes.NewModuleAddressOrBech32Address(in.Config.Authority)
 	}
 
-	k := keeper.NewKeeper(in.Codec, in.StoreService, in.AccountKeeper, in.BankKeeper, in.StakingKeeper, authority.String())
+	k := keeper.NewKeeper(in.Codec, in.Environment, in.AccountKeeper, in.BankKeeper, in.StakingKeeper, authority.String())
 	m := NewAppModule(in.Codec, k, in.AccountKeeper, in.BankKeeper)
 
 	return ModuleOutputs{

x/protocolpool/keeper/genesis.go

@@ -10,8 +10,7 @@ import (
 )
 
 func (k Keeper) InitGenesis(ctx context.Context, data *types.GenesisState) error {
-	sdkCtx := sdk.UnwrapSDKContext(ctx)
-	currentTime := sdkCtx.BlockTime()
+	currentTime := k.environment.HeaderService.GetHeaderInfo(ctx).Time
 	for _, cf := range data.ContinuousFund {
 		// ignore expired ContinuousFunds
 		if cf.Expiry != nil && cf.Expiry.Before(currentTime) {

x/protocolpool/keeper/grpc_query_test.go

@@ -10,7 +10,7 @@ import (
 )
 
 func (suite *KeeperTestSuite) TestUnclaimedBudget() {
-	startTime := suite.ctx.BlockTime().Add(-70 * time.Second)
+	startTime := suite.environment.HeaderService.GetHeaderInfo(suite.ctx).Time.Add(-70 * time.Second)
 	period := time.Duration(60) * time.Second
 	zeroCoin := sdk.NewCoin("foo", math.ZeroInt())
 	nextClaimFrom := startTime.Add(period)

x/protocolpool/keeper/keeper.go

@@ -8,7 +8,7 @@ import (
 	"time"
 
 	"cosmossdk.io/collections"
-	storetypes "cosmossdk.io/core/store"
+	"cosmossdk.io/core/appmodule"
 	errorsmod "cosmossdk.io/errors"
 	"cosmossdk.io/log"
 	"cosmossdk.io/math"
@@ -20,7 +20,7 @@ import (
 )
 
 type Keeper struct {
-	storeService  storetypes.KVStoreService
+	environment   appmodule.Environment
 	authKeeper    types.AccountKeeper
 	bankKeeper    types.BankKeeper
 	stakingKeeper types.StakingKeeper
@@ -41,8 +41,7 @@ type Keeper struct {
 	ToDistribute collections.Item[math.Int]
 }
 
-func NewKeeper(cdc codec.BinaryCodec, storeService storetypes.KVStoreService,
-	ak types.AccountKeeper, bk types.BankKeeper, sk types.StakingKeeper, authority string,
+func NewKeeper(cdc codec.BinaryCodec, env appmodule.Environment, ak types.AccountKeeper, bk types.BankKeeper, sk types.StakingKeeper, authority string,
 ) Keeper {
 	// ensure pool module account is set
 	if addr := ak.GetModuleAddress(types.ModuleName); addr == nil {
@@ -53,10 +52,10 @@ func NewKeeper(cdc codec.BinaryCodec, storeService storetypes.KVStoreService,
 		panic(fmt.Sprintf("%s module account has not been set", types.StreamAccount))
 	}
 
-	sb := collections.NewSchemaBuilder(storeService)
+	sb := collections.NewSchemaBuilder(env.KVStoreService)
 
 	keeper := Keeper{
-		storeService:              storeService,
+		environment:               env,
 		authKeeper:                ak,
 		bankKeeper:                bk,
 		stakingKeeper:             sk,
@@ -85,8 +84,7 @@ func (k Keeper) GetAuthority() string {
 
 // Logger returns a module-specific logger.
 func (k Keeper) Logger(ctx context.Context) log.Logger {
-	sdkCtx := sdk.UnwrapSDKContext(ctx)
-	return sdkCtx.Logger().With(log.ModuleKey, "x/"+types.ModuleName)
+	return k.environment.Logger.With(log.ModuleKey, "x/"+types.ModuleName)
 }
 
 // FundCommunityPool allows an account to directly fund the community fund pool.
@@ -116,15 +114,14 @@ func (k Keeper) GetCommunityPool(ctx context.Context) (sdk.Coins, error) {
 }
 
 func (k Keeper) withdrawContinuousFund(ctx context.Context, recipient sdk.AccAddress) (sdk.Coin, error) {
-	sdkCtx := sdk.UnwrapSDKContext(ctx)
 	cf, err := k.ContinuousFund.Get(ctx, recipient)
 	if err != nil {
 		if errors.Is(err, collections.ErrNotFound) {
 			return sdk.Coin{}, fmt.Errorf("no continuous fund found for recipient: %s", recipient.String())
 		}
 		return sdk.Coin{}, fmt.Errorf("get continuous fund failed for recipient: %s", recipient.String())
 	}
-	if cf.Expiry != nil && cf.Expiry.Before(sdkCtx.HeaderInfo().Time) {
+	if cf.Expiry != nil && cf.Expiry.Before(k.environment.HeaderService.GetHeaderInfo(ctx).Time) {
 		return sdk.Coin{}, fmt.Errorf("cannot withdraw continuous funds: continuous fund expired for recipient: %s", recipient.String())
 	}
 
@@ -326,8 +323,6 @@ func (k Keeper) claimFunds(ctx context.Context, recipient sdk.AccAddress) (amoun
 }
 
 func (k Keeper) getClaimableFunds(ctx context.Context, recipient sdk.AccAddress) (amount sdk.Coin, err error) {
-	sdkCtx := sdk.UnwrapSDKContext(ctx)
-
 	budget, err := k.BudgetProposal.Get(ctx, recipient)
 	if err != nil {
 		if errors.Is(err, collections.ErrNotFound) {
@@ -349,7 +344,7 @@ func (k Keeper) getClaimableFunds(ctx context.Context, recipient sdk.AccAddress)
 		}
 	}
 
-	currentTime := sdkCtx.BlockTime()
+	currentTime := k.environment.HeaderService.GetHeaderInfo(ctx).Time
 	startTime := budget.StartTime
 
 	// Check if the start time is reached
@@ -416,7 +411,7 @@ func (k Keeper) validateAndUpdateBudgetProposal(ctx context.Context, bp types.Ms
 		return nil, fmt.Errorf("invalid budget proposal: %w", err)
 	}
 
-	currentTime := sdk.UnwrapSDKContext(ctx).BlockTime()
+	currentTime := k.environment.HeaderService.GetHeaderInfo(ctx).Time
 	if bp.StartTime.IsZero() || bp.StartTime == nil {
 		bp.StartTime = &currentTime
 	}
@@ -459,7 +454,7 @@ func (k Keeper) validateContinuousFund(ctx context.Context, msg types.MsgCreateC
 	}
 
 	// Validate expiry
-	currentTime := sdk.UnwrapSDKContext(ctx).BlockTime()
+	currentTime := k.environment.HeaderService.GetHeaderInfo(ctx).Time
 	if msg.Expiry != nil && msg.Expiry.Compare(currentTime) == -1 {
 		return fmt.Errorf("expiry time cannot be less than the current block time")
 	}

x/protocolpool/keeper/keeper_test.go

@@ -7,7 +7,9 @@ import (
 	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/suite"
 
+	"cosmossdk.io/core/appmodule"
 	"cosmossdk.io/core/header"
+	"cosmossdk.io/log"
 	"cosmossdk.io/math"
 	storetypes "cosmossdk.io/store/types"
 	authtypes "cosmossdk.io/x/auth/types"
@@ -32,6 +34,7 @@ type KeeperTestSuite struct {
 	suite.Suite
 
 	ctx           sdk.Context
+	environment   appmodule.Environment
 	poolKeeper    poolkeeper.Keeper
 	authKeeper    *pooltestutil.MockAccountKeeper
 	bankKeeper    *pooltestutil.MockBankKeeper
@@ -44,6 +47,7 @@ type KeeperTestSuite struct {
 func (s *KeeperTestSuite) SetupTest() {
 	key := storetypes.NewKVStoreKey(types.StoreKey)
 	storeService := runtime.NewKVStoreService(key)
+	environment := runtime.NewEnvironment(storeService, log.NewNopLogger())
 	testCtx := testutil.DefaultContextWithDB(s.T(), key, storetypes.NewTransientStoreKey("transient_test"))
 	ctx := testCtx.Ctx.WithHeaderInfo(header.Info{Time: time.Now()})
 	encCfg := moduletestutil.MakeTestEncodingConfig()
@@ -65,14 +69,15 @@ func (s *KeeperTestSuite) SetupTest() {
 
 	poolKeeper := poolkeeper.NewKeeper(
 		encCfg.Codec,
-		storeService,
+		environment,
 		accountKeeper,
 		bankKeeper,
 		stakingKeeper,
 		authtypes.NewModuleAddress(types.GovModuleName).String(),
 	)
 	s.ctx = ctx
 	s.poolKeeper = poolKeeper
+	s.environment = environment
 
 	types.RegisterInterfaces(encCfg.InterfaceRegistry)
 	queryHelper := baseapp.NewQueryServerTestHelper(ctx, encCfg.InterfaceRegistry)

x/protocolpool/keeper/msg_server.go

@@ -181,8 +181,6 @@ func (k MsgServer) WithdrawContinuousFund(ctx context.Context, msg *types.MsgWit
 }
 
 func (k MsgServer) CancelContinuousFund(ctx context.Context, msg *types.MsgCancelContinuousFund) (*types.MsgCancelContinuousFundResponse, error) {
-	sdkCtx := sdk.UnwrapSDKContext(ctx)
-
 	if err := k.validateAuthority(msg.Authority); err != nil {
 		return nil, err
 	}
@@ -192,8 +190,8 @@ func (k MsgServer) CancelContinuousFund(ctx context.Context, msg *types.MsgCance
 		return nil, err
 	}
 
-	canceledHeight := sdkCtx.BlockHeight()
-	canceledTime := sdkCtx.BlockTime()
+	canceledHeight := k.environment.HeaderService.GetHeaderInfo(ctx).Height
+	canceledTime := k.environment.HeaderService.GetHeaderInfo(ctx).Time
 
 	found, err := k.ContinuousFund.Has(ctx, recipient)
 	if !found {