From 50f222bc3004da70748b4c05db811bf7981ead40 Mon Sep 17 00:00:00 2001
From: lazar <lukic.lazar@yahoo.com>
Date: Thu, 3 Oct 2024 17:09:57 +0200
Subject: [PATCH] pr comments

---
 .../expected_babylon_client.go                |  3 ++-
 .../mock_babylon_client.go                    | 16 +++++++++++++++
 .../stakingeventwatcher.go                    | 20 +++++++++++--------
 3 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/btcstaking-tracker/stakingeventwatcher/expected_babylon_client.go b/btcstaking-tracker/stakingeventwatcher/expected_babylon_client.go
index befd5685..586c6e30 100644
--- a/btcstaking-tracker/stakingeventwatcher/expected_babylon_client.go
+++ b/btcstaking-tracker/stakingeventwatcher/expected_babylon_client.go
@@ -25,6 +25,7 @@ type Delegation struct {
 
 type BabylonNodeAdapter interface {
 	BtcDelegations(offset uint64, limit uint64) ([]Delegation, error)
+	DelegationsByStatus(status btcstakingtypes.BTCDelegationStatus, offset uint64, limit uint64) ([]Delegation, error)
 	IsDelegationActive(stakingTxHash chainhash.Hash) (bool, error)
 	IsDelegationVerified(stakingTxHash chainhash.Hash) (bool, error)
 	ReportUnbonding(ctx context.Context, stakingTxHash chainhash.Hash, stakerUnbondingSig *schnorr.Signature) error
@@ -175,7 +176,7 @@ func (bca *BabylonClientAdapter) ActivateDelegation(
 	resp, err := bca.babylonClient.ReliablySendMsg(ctx, &msg, []*errors.Error{}, []*errors.Error{})
 
 	if err != nil && resp != nil {
-		return fmt.Errorf("msg MsgBTCUndelegate failed exeuction with code %d and error %w", resp.Code, err)
+		return fmt.Errorf("msg MsgAddBTCDelegationInclusionProof failed exeuction with code %d and error %w", resp.Code, err)
 	}
 
 	if err != nil {
diff --git a/btcstaking-tracker/stakingeventwatcher/mock_babylon_client.go b/btcstaking-tracker/stakingeventwatcher/mock_babylon_client.go
index e362ece2..fcb0a4ed 100644
--- a/btcstaking-tracker/stakingeventwatcher/mock_babylon_client.go
+++ b/btcstaking-tracker/stakingeventwatcher/mock_babylon_client.go
@@ -9,6 +9,7 @@ import (
 	reflect "reflect"
 
 	types "github.com/babylonlabs-io/babylon/x/btccheckpoint/types"
+	types0 "github.com/babylonlabs-io/babylon/x/btcstaking/types"
 	schnorr "github.com/btcsuite/btcd/btcec/v2/schnorr"
 	chainhash "github.com/btcsuite/btcd/chaincfg/chainhash"
 	gomock "github.com/golang/mock/gomock"
@@ -81,6 +82,21 @@ func (mr *MockBabylonNodeAdapterMockRecorder) BtcDelegations(offset, limit inter
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "BtcDelegations", reflect.TypeOf((*MockBabylonNodeAdapter)(nil).BtcDelegations), offset, limit)
 }
 
+// DelegationsByStatus mocks base method.
+func (m *MockBabylonNodeAdapter) DelegationsByStatus(status types0.BTCDelegationStatus, offset, limit uint64) ([]Delegation, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "DelegationsByStatus", status, offset, limit)
+	ret0, _ := ret[0].([]Delegation)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// DelegationsByStatus indicates an expected call of DelegationsByStatus.
+func (mr *MockBabylonNodeAdapterMockRecorder) DelegationsByStatus(status, offset, limit interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DelegationsByStatus", reflect.TypeOf((*MockBabylonNodeAdapter)(nil).DelegationsByStatus), status, offset, limit)
+}
+
 // IsDelegationActive mocks base method.
 func (m *MockBabylonNodeAdapter) IsDelegationActive(stakingTxHash chainhash.Hash) (bool, error) {
 	m.ctrl.T.Helper()
diff --git a/btcstaking-tracker/stakingeventwatcher/stakingeventwatcher.go b/btcstaking-tracker/stakingeventwatcher/stakingeventwatcher.go
index 05c46298..8a169386 100644
--- a/btcstaking-tracker/stakingeventwatcher/stakingeventwatcher.go
+++ b/btcstaking-tracker/stakingeventwatcher/stakingeventwatcher.go
@@ -127,7 +127,7 @@ func (sew *StakingEventWatcher) Start() error {
 
 		sew.wg.Add(4)
 		go sew.handleNewBlocks(blockEventNotifier)
-		go sew.handleDelegations()
+		go sew.handleUnbondedDelegations()
 		go sew.fetchDelegations()
 		go sew.handlerVerifiedDelegations()
 
@@ -189,7 +189,8 @@ func (sew *StakingEventWatcher) checkBabylonDelegations() error {
 			}
 
 			// if we already have this delegation, skip it
-			if sew.unbondingTracker.GetDelegation(stakingTxHash) == nil && delegation.HasProof {
+			// we should track both verified and active status for unbonding
+			if sew.unbondingTracker.GetDelegation(stakingTxHash) == nil {
 				utils.PushOrQuit(sew.unbondingDelegationChan, del, sew.quit)
 			}
 
@@ -334,15 +335,18 @@ func (sew *StakingEventWatcher) reportUnbondingToBabylon(
 			return fmt.Errorf("error checking if delegation is active: %v", err)
 		}
 
-		if !active {
-			//
+		verified, err := sew.babylonNodeAdapter.IsDelegationVerified(stakingTxHash)
+
+		if err != nil {
+			return fmt.Errorf("error checking if delegation is verified: %v", err)
+		}
+
+		if !active && !verified {
 			sew.logger.Debugf("cannot report unbonding. delegation for staking tx %s is no longer active", stakingTxHash)
 			return nil
 		}
 
-		err = sew.babylonNodeAdapter.ReportUnbonding(ctx, stakingTxHash, unbondingSignature)
-
-		if err != nil {
+		if err = sew.babylonNodeAdapter.ReportUnbonding(ctx, stakingTxHash, unbondingSignature); err != nil {
 			sew.metrics.FailedReportedUnbondingTransactions.Inc()
 			return fmt.Errorf("error reporting unbonding tx %s to babylon: %v", stakingTxHash, err)
 		}
@@ -402,7 +406,7 @@ func (sew *StakingEventWatcher) watchForSpend(spendEvent *notifier.SpendEvent, t
 	)
 }
 
-func (sew *StakingEventWatcher) handleDelegations() {
+func (sew *StakingEventWatcher) handleUnbondedDelegations() {
 	defer sew.wg.Done()
 	for {
 		select {