From e65160b893975a03cd80d40e221a111d26418cff Mon Sep 17 00:00:00 2001
From: KonradStaniec <konrad.staniec@gmail.com>
Date: Thu, 5 Dec 2024 08:59:13 +0100
Subject: [PATCH] pr comments

---
 babylonclient/babyloncontroller.go | 72 +++++++++++++++---------------
 babylonclient/interface.go         |  7 +--
 staker/stakerapp.go                |  1 -
 3 files changed, 37 insertions(+), 43 deletions(-)

diff --git a/babylonclient/babyloncontroller.go b/babylonclient/babyloncontroller.go
index f26055d..406c5ca 100644
--- a/babylonclient/babyloncontroller.go
+++ b/babylonclient/babyloncontroller.go
@@ -164,13 +164,7 @@ func (bc *BabylonController) BTCCheckpointParams() (*BTCCheckpointParams, error)
 	return bc.btccheckpointParamsWithRetry()
 }
 
-func (bc *BabylonController) Params() (*StakingParams, error) {
-	bccParams, err := bc.btccheckpointParamsWithRetry()
-
-	if err != nil {
-		return nil, err
-	}
-
+func (bc *BabylonController) queryStakingTrackerWithRetries() (*StakingTrackerResponse, error) {
 	var stakingTrackerParams *StakingTrackerResponse
 	if err := retry.Do(func() error {
 		trackerParams, err := bc.QueryStakingTracker()
@@ -189,6 +183,22 @@ func (bc *BabylonController) Params() (*StakingParams, error) {
 		return nil, err
 	}
 
+	return stakingTrackerParams, nil
+}
+
+func (bc *BabylonController) Params() (*StakingParams, error) {
+	bccParams, err := bc.btccheckpointParamsWithRetry()
+
+	if err != nil {
+		return nil, err
+	}
+
+	stakingTrackerParams, err := bc.queryStakingTrackerWithRetries()
+
+	if err != nil {
+		return nil, err
+	}
+
 	return &StakingParams{
 		ConfirmationTimeBlocks:    bccParams.ConfirmationTimeBlocks,
 		FinalizationTimeoutBlocks: bccParams.FinalizationTimeoutBlocks,
@@ -207,13 +217,9 @@ func (bc *BabylonController) Params() (*StakingParams, error) {
 	}, nil
 }
 
-func (bc *BabylonController) ParamsByBtcHeight(btcHeight uint32) (*StakingParams, error) {
-	bccParams, err := bc.btccheckpointParamsWithRetry()
-
-	if err != nil {
-		return nil, err
-	}
-
+func (bc *BabylonController) queryStakingTrackerByBtcHeightWithRetries(
+	btcHeight uint32,
+) (*StakingTrackerResponse, error) {
 	var stakingTrackerParams *StakingTrackerResponse
 	if err := retry.Do(func() error {
 		trackerParams, err := bc.QueryStakingTrackerByBtcHeight(btcHeight)
@@ -232,6 +238,22 @@ func (bc *BabylonController) ParamsByBtcHeight(btcHeight uint32) (*StakingParams
 		return nil, err
 	}
 
+	return stakingTrackerParams, nil
+}
+
+func (bc *BabylonController) ParamsByBtcHeight(btcHeight uint32) (*StakingParams, error) {
+	bccParams, err := bc.btccheckpointParamsWithRetry()
+
+	if err != nil {
+		return nil, err
+	}
+
+	stakingTrackerParams, err := bc.queryStakingTrackerByBtcHeightWithRetries(btcHeight)
+
+	if err != nil {
+		return nil, err
+	}
+
 	return &StakingParams{
 		ConfirmationTimeBlocks:    bccParams.ConfirmationTimeBlocks,
 		FinalizationTimeoutBlocks: bccParams.FinalizationTimeoutBlocks,
@@ -250,28 +272,6 @@ func (bc *BabylonController) ParamsByBtcHeight(btcHeight uint32) (*StakingParams
 	}, nil
 }
 
-func (bc *BabylonController) StakingTrackerByBtcHeight(btcHeight uint32) (*StakingTrackerResponse, error) {
-	var stakingTrackerParams *StakingTrackerResponse
-	if err := retry.Do(func() error {
-		trackerParams, err := bc.QueryStakingTrackerByBtcHeight(btcHeight)
-		if err != nil {
-			return err
-		}
-		stakingTrackerParams = trackerParams
-		return nil
-	}, RtyAtt, RtyDel, RtyErr, retry.OnRetry(func(n uint, err error) {
-		bc.logger.WithFields(logrus.Fields{
-			"attempt":      n + 1,
-			"max_attempts": RtyAttNum,
-			"error":        err,
-		}).Error("Failed to query babylon client for staking tracker params")
-	})); err != nil {
-		return nil, err
-	}
-
-	return stakingTrackerParams, nil
-}
-
 func (bc *BabylonController) GetKeyAddress() sdk.AccAddress {
 	// get key address, retrieves address based on key name which is configured in
 	// cfg *stakercfg.BBNConfig. If this fails, it means we have misconfiguration problem
diff --git a/babylonclient/interface.go b/babylonclient/interface.go
index 154464a..f045b5c 100644
--- a/babylonclient/interface.go
+++ b/babylonclient/interface.go
@@ -78,7 +78,6 @@ type BabylonClient interface {
 	BTCCheckpointParams() (*BTCCheckpointParams, error)
 	Params() (*StakingParams, error)
 	ParamsByBtcHeight(btcHeight uint32) (*StakingParams, error)
-	StakingTrackerByBtcHeight(btcHeight uint32) (*StakingTrackerResponse, error)
 	Delegate(dg *DelegationData) (*pv.RelayerTxResponse, error)
 	QueryFinalityProviders(limit uint64, offset uint64) (*FinalityProvidersClientResponse, error)
 	QueryFinalityProvider(btcPubKey *btcec.PublicKey) (*FinalityProviderClientResponse, error)
@@ -101,7 +100,7 @@ func (m *MockBabylonClient) Params() (*StakingParams, error) {
 	return m.ClientParams, nil
 }
 
-func (m *MockBabylonClient) ParamsByBtcHeight(btcHeight uint32) (*StakingParams, error) {
+func (m *MockBabylonClient) ParamsByBtcHeight(_ uint32) (*StakingParams, error) {
 	return m.ClientParams, nil
 }
 
@@ -112,10 +111,6 @@ func (m *MockBabylonClient) BTCCheckpointParams() (*BTCCheckpointParams, error)
 	}, nil
 }
 
-func (m *MockBabylonClient) StakingTrackerByBtcHeight(btcHeight uint32) (*StakingTrackerResponse, error) {
-	return &StakingTrackerResponse{}, nil
-}
-
 func (m *MockBabylonClient) Sign(msg []byte) ([]byte, error) {
 	sig, err := m.babylonKey.Sign(msg)
 
diff --git a/staker/stakerapp.go b/staker/stakerapp.go
index 4ecebe1..9e6b70d 100644
--- a/staker/stakerapp.go
+++ b/staker/stakerapp.go
@@ -984,7 +984,6 @@ func (app *App) retrieveExternalDelegationData(
 		}
 
 		params = p
-
 	} else {
 		p, err := app.babylonClient.ParamsByBtcHeight(inclusionInfo.inclusionBlockBtcHeight)