chore(e2e): e2e test for activating and unbonded in mempool (#71)

This PR introduces an end-to-end test to validate that when both a
staking transaction and an unbonding transaction are included in the
same block, the delegation correctly transitions to the "UNBONDED"
status.

btcstaking-tracker/stakingeventwatcher/stakingeventwatcher.go

@@ -207,7 +207,7 @@ func (sew *StakingEventWatcher) fetchDelegations() {
 				continue
 			}
 
-			addToUnbonding := func(delegation Delegation) {
+			addToUnbondingFunc := func(delegation Delegation) {
 				del := &newDelegation{
 					stakingTxHash:         delegation.StakingTx.TxHash(),
 					stakingTx:             delegation.StakingTx,
@@ -216,14 +216,19 @@ func (sew *StakingEventWatcher) fetchDelegations() {
 					unbondingOutput:       delegation.UnbondingOutput,
 				}
 
-				// if we already have this delegation, skip it
+				// if we already have this delegation, we still want to check if it has changed,
 				// we should track both verified and active status for unbonding
-				if sew.unbondingTracker.GetDelegation(delegation.StakingTx.TxHash()) == nil {
+				changed, exists := sew.unbondingTracker.HasDelegationChanged(delegation.StakingTx.TxHash(), del)
+				if exists && changed {
+					// Delegation exists and has changed, push the update.
+					utils.PushOrQuit(sew.unbondingDelegationChan, del, sew.quit)
+				} else if !exists {
+					// Delegation doesn't exist, push the new delegation.
 					utils.PushOrQuit(sew.unbondingDelegationChan, del, sew.quit)
 				}
 			}
 
-			addToPending := func(delegation Delegation) {
+			addToPendingFunc := func(delegation Delegation) {
 				del := &newDelegation{
 					stakingTxHash:         delegation.StakingTx.TxHash(),
 					stakingTx:             delegation.StakingTx,
@@ -237,6 +242,8 @@ func (sew *StakingEventWatcher) fetchDelegations() {
 						del.stakingTx,
 						del.stakingOutputIdx,
 						del.unbondingOutput,
+						del.delegationStartHeight,
+						false,
 					)
 				}
 			}
@@ -246,21 +253,21 @@ func (sew *StakingEventWatcher) fetchDelegations() {
 
 			go func() {
 				defer wg.Done()
-				if err = sew.checkBabylonDelegations(btcstakingtypes.BTCDelegationStatus_ACTIVE, addToUnbonding); err != nil {
+				if err = sew.checkBabylonDelegations(btcstakingtypes.BTCDelegationStatus_ACTIVE, addToUnbondingFunc); err != nil {
 					sew.logger.Errorf("error checking babylon delegations: %v", err)
 				}
 			}()
 
 			go func() {
 				defer wg.Done()
-				if err = sew.checkBabylonDelegations(btcstakingtypes.BTCDelegationStatus_VERIFIED, addToUnbonding); err != nil {
+				if err = sew.checkBabylonDelegations(btcstakingtypes.BTCDelegationStatus_VERIFIED, addToUnbondingFunc); err != nil {
 					sew.logger.Errorf("error checking babylon delegations: %v", err)
 				}
 			}()
 
 			go func() {
 				defer wg.Done()
-				if err = sew.checkBabylonDelegations(btcstakingtypes.BTCDelegationStatus_VERIFIED, addToPending); err != nil {
+				if err = sew.checkBabylonDelegations(btcstakingtypes.BTCDelegationStatus_VERIFIED, addToPendingFunc); err != nil {
 					sew.logger.Errorf("error checking babylon delegations: %v", err)
 				}
 			}()
@@ -457,6 +464,8 @@ func (sew *StakingEventWatcher) handleUnbondedDelegations() {
 				activeDel.stakingTx,
 				activeDel.stakingOutputIdx,
 				activeDel.unbondingOutput,
+				activeDel.delegationStartHeight,
+				true,
 			)
 
 			if err != nil {

btcstaking-tracker/stakingeventwatcher/tracked_delegations.go

@@ -9,9 +9,10 @@ import (
 )
 
 type TrackedDelegation struct {
-	StakingTx        *wire.MsgTx
-	StakingOutputIdx uint32
-	UnbondingOutput  *wire.TxOut
+	StakingTx             *wire.MsgTx
+	StakingOutputIdx      uint32
+	UnbondingOutput       *wire.TxOut
+	DelegationStartHeight uint64
 }
 
 type TrackedDelegations struct {
@@ -60,11 +61,14 @@ func (td *TrackedDelegations) AddDelegation(
 	StakingTx *wire.MsgTx,
 	StakingOutputIdx uint32,
 	UnbondingOutput *wire.TxOut,
+	delegationStartHeight uint64,
+	shouldUpdate bool,
 ) (*TrackedDelegation, error) {
 	delegation := &TrackedDelegation{
-		StakingTx:        StakingTx,
-		StakingOutputIdx: StakingOutputIdx,
-		UnbondingOutput:  UnbondingOutput,
+		StakingTx:             StakingTx,
+		StakingOutputIdx:      StakingOutputIdx,
+		UnbondingOutput:       UnbondingOutput,
+		DelegationStartHeight: delegationStartHeight,
 	}
 
 	stakingTxHash := StakingTx.TxHash()
@@ -73,6 +77,11 @@ func (td *TrackedDelegations) AddDelegation(
 	defer td.mu.Unlock()
 
 	if _, ok := td.mapping[stakingTxHash]; ok {
+		if shouldUpdate {
+			// Update the existing delegation
+			td.mapping[stakingTxHash] = delegation
+			return delegation, nil
+		}
 		return nil, fmt.Errorf("delegation already tracked for staking tx hash %s", stakingTxHash)
 	}
 
@@ -86,3 +95,26 @@ func (td *TrackedDelegations) RemoveDelegation(stakingTxHash chainhash.Hash) {
 
 	delete(td.mapping, stakingTxHash)
 }
+
+func (td *TrackedDelegations) HasDelegationChanged(
+	stakingTxHash chainhash.Hash,
+	newDelegation *newDelegation,
+) (exists bool, changed bool) {
+	td.mu.Lock()
+	defer td.mu.Unlock()
+
+	// Check if the delegation exists in the map
+	existingDelegation, exists := td.mapping[stakingTxHash]
+	if !exists {
+		// If it doesn't exist, return false for changed, and false for exists
+		return false, false
+	}
+
+	// Compare height to check if the delegation has changed
+	if existingDelegation.DelegationStartHeight != newDelegation.delegationStartHeight {
+		return true, true // The delegation has changed and it exists
+	}
+
+	// The delegation exists but hasn't changed
+	return false, true
+}

e2etest/test_manager_btcstaking.go

@@ -214,7 +214,7 @@ func (tm *TestManager) CreateBTCDelegation(
 func (tm *TestManager) CreateBTCDelegationWithoutIncl(
 	t *testing.T,
 	fpSK *btcec.PrivateKey,
-) (*datagen.TestStakingSlashingInfo, *datagen.TestUnbondingSlashingInfo, *btcec.PrivateKey) {
+) (*wire.MsgTx, *datagen.TestStakingSlashingInfo, *datagen.TestUnbondingSlashingInfo, *btcec.PrivateKey) {
 	signerAddr := tm.BabylonClient.MustGetAddr()
 	addr := sdk.MustAccAddressFromBech32(signerAddr)
 
@@ -311,34 +311,7 @@ func (tm *TestManager) CreateBTCDelegationWithoutIncl(
 		stakingOutIdx,
 	)
 
-	// send staking tx to Bitcoin node's mempool
-	_, err = tm.BTCClient.SendRawTransaction(stakingMsgTx, true)
-	require.NoError(t, err)
-
-	require.Eventually(t, func() bool {
-		return len(tm.RetrieveTransactionFromMempool(t, []*chainhash.Hash{stakingMsgTxHash})) == 1
-	}, eventuallyWaitTimeOut, eventuallyPollTime)
-
-	mBlock := tm.mineBlock(t)
-	require.Equal(t, 2, len(mBlock.Transactions))
-
-	// wait until staking tx is on Bitcoin
-	require.Eventually(t, func() bool {
-		_, err := tm.BTCClient.GetRawTransaction(stakingMsgTxHash)
-		return err == nil
-	}, eventuallyWaitTimeOut, eventuallyPollTime)
-
-	// insert k empty blocks to Bitcoin
-	btccParamsResp, err := tm.BabylonClient.BTCCheckpointParams()
-	require.NoError(t, err)
-	btccParams := btccParamsResp.Params
-	for i := 0; i < int(btccParams.BtcConfirmationDepth); i++ {
-		tm.mineBlock(t)
-	}
-
-	tm.CatchUpBTCLightClient(t)
-
-	return stakingSlashingInfo, unbondingSlashingInfo, tm.WalletPrivKey
+	return stakingMsgTx, stakingSlashingInfo, unbondingSlashingInfo, tm.WalletPrivKey
 }
 
 func (tm *TestManager) createStakingAndSlashingTx(

e2etest/unbondingwatcher_e2e_test.go

@@ -4,6 +4,7 @@
 package e2etest
 
 import (
+	btcstakingtypes "github.com/babylonlabs-io/babylon/x/btcstaking/types"
 	"go.uber.org/zap"
 	"testing"
 	"time"
@@ -157,7 +158,35 @@ func TestActivatingDelegation(t *testing.T) {
 	// set up a finality provider
 	_, fpSK := tm.CreateFinalityProvider(t)
 	// set up a BTC delegation
-	stakingSlashingInfo, _, _ := tm.CreateBTCDelegationWithoutIncl(t, fpSK)
+	stakingMsgTx, stakingSlashingInfo, _, _ := tm.CreateBTCDelegationWithoutIncl(t, fpSK)
+	stakingMsgTxHash := stakingMsgTx.TxHash()
+
+	// send staking tx to Bitcoin node's mempool
+	_, err = tm.BTCClient.SendRawTransaction(stakingMsgTx, true)
+	require.NoError(t, err)
+
+	require.Eventually(t, func() bool {
+		return len(tm.RetrieveTransactionFromMempool(t, []*chainhash.Hash{&stakingMsgTxHash})) == 1
+	}, eventuallyWaitTimeOut, eventuallyPollTime)
+
+	mBlock := tm.mineBlock(t)
+	require.Equal(t, 2, len(mBlock.Transactions))
+
+	// wait until staking tx is on Bitcoin
+	require.Eventually(t, func() bool {
+		_, err := tm.BTCClient.GetRawTransaction(&stakingMsgTxHash)
+		return err == nil
+	}, eventuallyWaitTimeOut, eventuallyPollTime)
+
+	// insert k empty blocks to Bitcoin
+	btccParamsResp, err := tm.BabylonClient.BTCCheckpointParams()
+	require.NoError(t, err)
+	btccParams := btccParamsResp.Params
+	for i := 0; i < int(btccParams.BtcConfirmationDepth); i++ {
+		tm.mineBlock(t)
+	}
+
+	tm.CatchUpBTCLightClient(t)
 
 	// created delegation lacks inclusion proof, once created it will be in
 	// pending status, once convenant signatures are added it will be in verified status,
@@ -170,3 +199,117 @@ func TestActivatingDelegation(t *testing.T) {
 		return resp.BtcDelegation.Active
 	}, eventuallyWaitTimeOut, eventuallyPollTime)
 }
+
+// TestActivatingAndUnbondingDelegation tests that delegation will eventually become UNBONDED given that
+// both staking and unbonding tx are in the same block.
+// In this test, we include both staking tx and unbonding tx in the same block.
+// The delegation goes through "VERIFIED" → "ACTIVE" → "UNBONDED" status throughout this test.
+func TestActivatingAndUnbondingDelegation(t *testing.T) {
+	//t.Parallel()
+	// segwit is activated at height 300. It's necessary for staking/slashing tx
+	numMatureOutputs := uint32(300)
+
+	tm := StartManager(t, numMatureOutputs, defaultEpochInterval)
+	defer tm.Stop(t)
+	// Insert all existing BTC headers to babylon node
+	tm.CatchUpBTCLightClient(t)
+
+	btcNotifier, err := btcclient.NewNodeBackend(
+		btcclient.ToBitcoindConfig(tm.Config.BTC),
+		&chaincfg.RegressionNetParams,
+		&btcclient.EmptyHintCache{},
+	)
+	require.NoError(t, err)
+
+	err = btcNotifier.Start()
+	require.NoError(t, err)
+
+	commonCfg := config.DefaultCommonConfig()
+	bstCfg := config.DefaultBTCStakingTrackerConfig()
+	bstCfg.CheckDelegationsInterval = 1 * time.Second
+	stakingTrackerMetrics := metrics.NewBTCStakingTrackerMetrics()
+
+	bsTracker := bst.NewBTCStakingTracker(
+		tm.BTCClient,
+		btcNotifier,
+		tm.BabylonClient,
+		&bstCfg,
+		&commonCfg,
+		zap.NewNop(),
+		stakingTrackerMetrics,
+	)
+	bsTracker.Start()
+	defer bsTracker.Stop()
+
+	// set up a finality provider
+	_, fpSK := tm.CreateFinalityProvider(t)
+	// set up a BTC delegation
+	stakingMsgTx, stakingSlashingInfo, unbondingSlashingInfo, delSK := tm.CreateBTCDelegationWithoutIncl(t, fpSK)
+	stakingMsgTxHash := stakingMsgTx.TxHash()
+
+	// send staking tx to Bitcoin node's mempool
+	_, err = tm.BTCClient.SendRawTransaction(stakingMsgTx, true)
+	require.NoError(t, err)
+
+	require.Eventually(t, func() bool {
+		return len(tm.RetrieveTransactionFromMempool(t, []*chainhash.Hash{&stakingMsgTxHash})) == 1
+	}, eventuallyWaitTimeOut, eventuallyPollTime)
+
+	// Staker unbonds by directly sending tx to btc network. Watcher should detect it and report to babylon.
+	unbondingPathSpendInfo, err := stakingSlashingInfo.StakingInfo.UnbondingPathSpendInfo()
+	require.NoError(t, err)
+	stakingOutIdx, err := outIdx(unbondingSlashingInfo.UnbondingTx, unbondingSlashingInfo.UnbondingInfo.UnbondingOutput)
+	require.NoError(t, err)
+
+	unbondingTxSchnorrSig, err := btcstaking.SignTxWithOneScriptSpendInputStrict(
+		unbondingSlashingInfo.UnbondingTx,
+		stakingSlashingInfo.StakingTx,
+		stakingOutIdx,
+		unbondingPathSpendInfo.GetPkScriptPath(),
+		delSK,
+	)
+	require.NoError(t, err)
+
+	resp, err := tm.BabylonClient.BTCDelegation(stakingSlashingInfo.StakingTx.TxHash().String())
+	require.NoError(t, err)
+
+	covenantSigs := resp.BtcDelegation.UndelegationResponse.CovenantUnbondingSigList
+	witness, err := unbondingPathSpendInfo.CreateUnbondingPathWitness(
+		[]*schnorr.Signature{covenantSigs[0].Sig.MustToBTCSig()},
+		unbondingTxSchnorrSig,
+	)
+	require.NoError(t, err)
+	unbondingSlashingInfo.UnbondingTx.TxIn[0].Witness = witness
+
+	// Send unbonding tx to Bitcoin
+	_, err = tm.BTCClient.SendRawTransaction(unbondingSlashingInfo.UnbondingTx, true)
+	require.NoError(t, err)
+
+	unbondingTxHash := unbondingSlashingInfo.UnbondingTx.TxHash()
+	t.Logf("submitted unbonding tx with hash %s", unbondingTxHash.String())
+	require.Eventually(t, func() bool {
+		return len(tm.RetrieveTransactionFromMempool(t, []*chainhash.Hash{&unbondingTxHash})) == 1
+	}, eventuallyWaitTimeOut, eventuallyPollTime)
+
+	mBlock := tm.mineBlock(t)
+	// both staking and unbonding txs are in this block
+	require.Equal(t, 3, len(mBlock.Transactions))
+
+	// insert k empty blocks to Bitcoin
+	btccParamsResp, err := tm.BabylonClient.BTCCheckpointParams()
+	require.NoError(t, err)
+	btccParams := btccParamsResp.Params
+	for i := 0; i < int(btccParams.BtcConfirmationDepth); i++ {
+		tm.mineBlock(t)
+	}
+
+	tm.CatchUpBTCLightClient(t)
+
+	// wait until delegation has become unbonded
+	require.Eventually(t, func() bool {
+		resp, err := tm.BabylonClient.BTCDelegation(stakingSlashingInfo.StakingTx.TxHash().String())
+		require.NoError(t, err)
+
+		return resp.BtcDelegation.StatusDesc == btcstakingtypes.BTCDelegationStatus_UNBONDED.String()
+	}, eventuallyWaitTimeOut, eventuallyPollTime)
+}