fix

tests/e2e/main_test.go

@@ -2,9 +2,7 @@ package e2e
 
 import (
 	"encoding/json"
-	"math/rand"
 	"testing"
-	"time"
 
 	"github.com/CosmWasm/wasmd/x/wasm/ibctesting"
 	"github.com/babylonchain/babylon-sdk/demo/app"
@@ -16,8 +14,6 @@ import (
 	"github.com/stretchr/testify/suite"
 )
 
-var r = rand.New(rand.NewSource(time.Now().Unix()))
-
 // In the Test function, we create and run the suite
 func TestBabylonSDKTestSuite(t *testing.T) {
 	suite.Run(t, new(BabylonSDKTestSuite))
@@ -128,16 +124,9 @@ func (s *BabylonSDKTestSuite) Test2MockConsumerFpDelegation() {
 	s.NotEmpty(consumerDels)
 }
 
-// TODO: trigger BeginBlock via s.ConsumerChain rather than ConsumerApp
-func (s *BabylonSDKTestSuite) Test3BeginBlock() {
-	err := s.ConsumerApp.BabylonKeeper.BeginBlocker(s.ConsumerChain.GetContext())
-	s.NoError(err)
-}
-
-// TODO: trigger EndBlock via s.ConsumerChain rather than ConsumerApp
-func (s *BabylonSDKTestSuite) Test4EndBlock() {
-	_, err := s.ConsumerApp.BabylonKeeper.EndBlocker(s.ConsumerChain.GetContext())
-	s.NoError(err)
+func (s *BabylonSDKTestSuite) Test3NextBlock() {
+	// this triggers BeginBlock and EndBlock
+	s.ConsumerChain.NextBlock()
 }
 
 // TearDownSuite runs once after all the suite's tests have been run

x/babylon/keeper/abci.go

@@ -1,22 +1,22 @@
 package keeper
 
 import (
+	"context"
 	"time"
 
 	"github.com/babylonchain/babylon-sdk/x/babylon/types"
 	abci "github.com/cometbft/cometbft/abci/types"
 	"github.com/cosmos/cosmos-sdk/telemetry"
-	sdk "github.com/cosmos/cosmos-sdk/types"
 )
 
-func (k *Keeper) BeginBlocker(ctx sdk.Context) error {
+func (k *Keeper) BeginBlocker(ctx context.Context) error {
 	defer telemetry.ModuleMeasureSince(types.ModuleName, time.Now(), telemetry.MetricKeyBeginBlocker)
 
 	return k.SendBeginBlockMsg(ctx)
 }
 
 // EndBlocker is called after every block
-func (k *Keeper) EndBlocker(ctx sdk.Context) ([]abci.ValidatorUpdate, error) {
+func (k *Keeper) EndBlocker(ctx context.Context) ([]abci.ValidatorUpdate, error) {
 	defer telemetry.ModuleMeasureSince(types.ModuleName, time.Now(), telemetry.MetricKeyEndBlocker)
 
 	if err := k.SendEndBlockMsg(ctx); err != nil {

x/babylon/keeper/wasm.go

@@ -1,6 +1,7 @@
 package keeper
 
 import (
+	"context"
 	"encoding/hex"
 	"encoding/json"
 
@@ -10,7 +11,9 @@ import (
 )
 
 // SendBeginBlockMsg sends a BeginBlock sudo message to the BTC staking contract via sudo
-func (k Keeper) SendBeginBlockMsg(ctx sdk.Context) error {
+func (k Keeper) SendBeginBlockMsg(c context.Context) error {
+	ctx := sdk.UnwrapSDKContext(c)
+
 	// get address of the BTC staking contract
 	addrStr := k.GetParams(ctx).BtcStakingContractAddress
 	if len(addrStr) == 0 {
@@ -33,7 +36,9 @@ func (k Keeper) SendBeginBlockMsg(ctx sdk.Context) error {
 }
 
 // SendEndBlockMsg sends a EndBlock sudo message to the BTC staking contract via sudo
-func (k Keeper) SendEndBlockMsg(ctx sdk.Context) error {
+func (k Keeper) SendEndBlockMsg(c context.Context) error {
+	ctx := sdk.UnwrapSDKContext(c)
+
 	// get address of the BTC staking contract
 	addrStr := k.GetParams(ctx).BtcStakingContractAddress
 	if len(addrStr) == 0 {

x/babylon/module.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 
+	"cosmossdk.io/core/appmodule"
 	abci "github.com/cometbft/cometbft/abci/types"
 	"github.com/grpc-ecosystem/grpc-gateway/runtime"
 	"github.com/spf13/cobra"
@@ -24,8 +25,10 @@ import (
 const ConsensusVersion = 1
 
 var (
-	_ module.AppModuleBasic = AppModuleBasic{}
-	_ module.AppModule      = AppModule{}
+	_ appmodule.AppModule       = AppModule{}
+	_ appmodule.HasBeginBlocker = AppModule{}
+	_ module.HasABCIEndBlock    = AppModule{}
+	_ module.AppModuleBasic     = AppModuleBasic{}
 )
 
 // AppModuleBasic defines the basic application module used by the babylon module.
@@ -136,12 +139,12 @@ func (AppModule) ConsensusVersion() uint64 {
 }
 
 // BeginBlock executed before every block
-func (am AppModule) BeginBlock(ctx sdk.Context) error {
+func (am AppModule) BeginBlock(ctx context.Context) error {
 	return am.k.BeginBlocker(ctx)
 }
 
 // EndBlock executed after every block. It returns no validator updates.
-func (am AppModule) EndBlock(ctx sdk.Context) ([]abci.ValidatorUpdate, error) {
+func (am AppModule) EndBlock(ctx context.Context) ([]abci.ValidatorUpdate, error) {
 	return am.k.EndBlocker(ctx)
 }