fix BeginBlocker/EndBlocker and e2e (#18)

The previous implementation does not really trigger Begin/EndBlocker due
to the mismatch of interfaces with Cosmos SDK. This PR fixes the issue,
adds more asserts and tests it in e2e.

Test plan: `make test-e2e`

TODO in subsequent PRs:

- e2e for testing finality round

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,7 +14,7 @@ import (
 	"github.com/stretchr/testify/suite"
 )
 
-var r = rand.New(rand.NewSource(time.Now().Unix()))
+var testMsg types.ExecuteMessage
 
 // In the Test function, we create and run the suite
 func TestBabylonSDKTestSuite(t *testing.T) {
@@ -109,8 +107,8 @@ func (s *BabylonSDKTestSuite) Test1ContractDeployment() {
 
 // TestExample is an example test case
 func (s *BabylonSDKTestSuite) Test2MockConsumerFpDelegation() {
-	msg := types.GenExecMessage()
-	msgBytes, err := json.Marshal(msg)
+	testMsg = types.GenExecMessage()
+	msgBytes, err := json.Marshal(testMsg)
 	s.NoError(err)
 
 	// send msg to BTC staking contract via admin account
@@ -126,20 +124,48 @@ func (s *BabylonSDKTestSuite) Test2MockConsumerFpDelegation() {
 	consumerDels, err := s.ConsumerCli.Query(s.ConsumerContract.BTCStaking, Query{"delegations": {}})
 	s.NoError(err)
 	s.NotEmpty(consumerDels)
+
+	// ensure the BTC staking is activated
+	resp, err := s.ConsumerCli.Query(s.ConsumerContract.BTCStaking, Query{"activated_height": {}})
+	s.NoError(err)
+	parsedActivatedHeight := resp["height"].(float64)
+	currentHeight := s.ConsumerChain.GetContext().BlockHeight()
+	s.Equal(uint64(parsedActivatedHeight), uint64(currentHeight))
 }
 
-// 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) Test5NextBlock() {
+	// get current height
+	height := s.ConsumerChain.GetContext().BlockHeight()
+	// ensure the current block is not indexed yet
+	_, err := s.ConsumerCli.Query(s.ConsumerContract.BTCStaking, Query{
+		"block": {
+			"height": uint64(height),
+		},
+	})
+	s.Error(err)
+
+	// this triggers BeginBlock and EndBlock
+	s.ConsumerChain.NextBlock()
+
+	// ensure the current block is indexed
+	_, err = s.ConsumerCli.Query(s.ConsumerContract.BTCStaking, Query{
+		"block": {
+			"height": uint64(height),
+		},
+	})
+	s.NoError(err)
+}
+
 // TearDownSuite runs once after all the suite's tests have been run
 func (s *BabylonSDKTestSuite) TearDownSuite() {
 	// Cleanup code here

tests/e2e/types/datagen.go

@@ -9,11 +9,9 @@ import (
 	sdkmath "cosmossdk.io/math"
 	"github.com/babylonchain/babylon/testutil/datagen"
 	bbn "github.com/babylonchain/babylon/types"
-	"github.com/babylonchain/babylon/x/btcstaking/types"
+	bstypes "github.com/babylonchain/babylon/x/btcstaking/types"
 	"github.com/btcsuite/btcd/chaincfg"
 	"github.com/stretchr/testify/require"
-
-	bstypes "github.com/babylonchain/babylon/x/btcstaking/types"
 )
 
 func GenExecMessage() ExecuteMessage {
@@ -53,8 +51,8 @@ func GenExecMessage() ExecuteMessage {
 	return executeMessage
 }
 
-func genBTCDelegation() (*types.Params, ActiveBtcDelegation) {
-	var net = &chaincfg.RegressionNetParams
+func genBTCDelegation() (*bstypes.Params, ActiveBtcDelegation) {
+	net := &chaincfg.RegressionNetParams
 	r := rand.New(rand.NewSource(time.Now().Unix()))
 	t := &testing.T{}
 
@@ -81,7 +79,7 @@ func genBTCDelegation() (*types.Params, ActiveBtcDelegation) {
 	unbondingTime := uint16(100) + 1
 	slashingChangeLockTime := unbondingTime
 
-	bsParams := &types.Params{
+	bsParams := &bstypes.Params{
 		CovenantPks:     bbn.NewBIP340PKsFromBTCPKs(covenantPKs),
 		CovenantQuorum:  covenantQuorum,
 		SlashingAddress: slashingAddress.EncodeAddress(),

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"
 
@@ -9,15 +10,40 @@ import (
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
 
-// SendBeginBlockMsg sends a BeginBlock sudo message to the BTC staking contract via sudo
-func (k Keeper) SendBeginBlockMsg(ctx sdk.Context) error {
+func (k Keeper) getBTCStakingContractAddr(ctx sdk.Context) sdk.AccAddress {
 	// get address of the BTC staking contract
 	addrStr := k.GetParams(ctx).BtcStakingContractAddress
 	if len(addrStr) == 0 {
 		// the BTC staking contract address is not set yet, skip sending BeginBlockMsg
 		return nil
 	}
-	addr := sdk.MustAccAddressFromBech32(addrStr)
+	addr, err := sdk.AccAddressFromBech32(addrStr)
+	if err != nil {
+		// Although this is a programming error so we should panic, we emit
+		// a warning message to minimise the impact on the consumer chain's operation
+		k.Logger(ctx).Warn("the BTC staking contract address is malformed", "contract", addrStr, "error", err)
+		return nil
+	}
+	if !k.wasm.HasContractInfo(ctx, addr) {
+		// NOTE: it's possible that the default contract address does not correspond to
+		// any contract. We emit a warning message rather than panic to minimise the
+		// impact on the consumer chain's operation
+		k.Logger(ctx).Warn("the BTC staking contract address is not on-chain", "contract", addrStr)
+		return nil
+	}
+
+	return addr
+}
+
+// SendBeginBlockMsg sends a BeginBlock sudo message to the BTC staking contract via sudo
+func (k Keeper) SendBeginBlockMsg(c context.Context) error {
+	ctx := sdk.UnwrapSDKContext(c)
+
+	// try to get and parse BTC staking contract
+	addr := k.getBTCStakingContractAddr(ctx)
+	if addr == nil {
+		return nil
+	}
 
 	// construct the sudo message
 	headerInfo := ctx.HeaderInfo()
@@ -33,14 +59,14 @@ 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 {
-	// get address of the BTC staking contract
-	addrStr := k.GetParams(ctx).BtcStakingContractAddress
-	if len(addrStr) == 0 {
-		// the BTC staking contract address is not set yet, skip sending EndBlockMsg
+func (k Keeper) SendEndBlockMsg(c context.Context) error {
+	ctx := sdk.UnwrapSDKContext(c)
+
+	// try to get and parse BTC staking contract
+	addr := k.getBTCStakingContractAddr(ctx)
+	if addr == nil {
 		return nil
 	}
-	addr := sdk.MustAccAddressFromBech32(addrStr)
 
 	// construct the sudo message
 	headerInfo := ctx.HeaderInfo()
@@ -61,6 +87,7 @@ func (k Keeper) doSudoCall(ctx sdk.Context, contractAddr sdk.AccAddress, msg con
 	if err != nil {
 		return errorsmod.Wrap(err, "marshal sudo msg")
 	}
-	_, err = k.wasm.Sudo(ctx, contractAddr, bz)
+	resp, err := k.wasm.Sudo(ctx, contractAddr, bz)
+	k.Logger(ctx).Debug("response of sudo call %v to contract %s: %v", bz, contractAddr.String(), resp)
 	return err
 }

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)
 }