From 4ba5e0c683f681d07554d125dc6aafbae4fc645b Mon Sep 17 00:00:00 2001
From: Crypto Minion <154598612+jrwbabylonlab@users.noreply.github.com>
Date: Tue, 26 Nov 2024 10:50:27 +1100
Subject: [PATCH] feat: return IsEligibleForTransition in delegation (#139)

* feat: return IsEligibleForTransition in delegation
---
 docs/docs.go                                  |   6 +-
 docs/swagger.json                             |   6 +-
 docs/swagger.yaml                             |   3 +
 internal/indexer/db/client/interface.go       |   6 +
 .../db/client/last_processed_height.go        |  25 ++++
 .../indexer/db/model/last_processed_height.go |   5 +
 internal/indexer/db/model/setup.go            |   1 +
 internal/shared/config/config.go              |  20 ++-
 .../shared/config/delegation_transition.go    |  26 ++++
 internal/shared/types/delegation.go           |   3 +
 internal/v1/api/handlers/delegation.go        |   3 +-
 internal/v1/db/model/delegation.go            |   2 +-
 internal/v1/queue/handler/expired_staking.go  |   3 +-
 internal/v1/queue/handler/unbonding.go        |   2 +-
 internal/v1/queue/handler/withdraw.go         |   2 +-
 internal/v1/service/delegation.go             | 126 ++++++++++++------
 internal/v1/service/interface.go              |   5 +-
 tests/mocks/mock_db_client.go                 |   2 +-
 tests/mocks/mock_ordinal_client.go            |   2 +-
 tests/mocks/mock_v1_db_client.go              |   2 +-
 tests/mocks/mock_v2_db_client.go              |  64 ++++++++-
 21 files changed, 250 insertions(+), 64 deletions(-)
 create mode 100644 internal/indexer/db/client/last_processed_height.go
 create mode 100644 internal/indexer/db/model/last_processed_height.go
 create mode 100644 internal/shared/config/delegation_transition.go

diff --git a/docs/docs.go b/docs/docs.go
index b4cc3235..41e6c5ed 100644
--- a/docs/docs.go
+++ b/docs/docs.go
@@ -196,7 +196,8 @@ const docTemplate = `{
                             "unbonding_requested",
                             "unbonding",
                             "unbonded",
-                            "withdrawn"
+                            "withdrawn",
+                            "transitioned"
                         ],
                         "type": "string",
                         "description": "Filter by state",
@@ -999,6 +1000,9 @@ const docTemplate = `{
                 "finality_provider_pk_hex": {
                     "type": "string"
                 },
+                "is_eligible_for_transition": {
+                    "type": "boolean"
+                },
                 "is_overflow": {
                     "type": "boolean"
                 },
diff --git a/docs/swagger.json b/docs/swagger.json
index dc4eece6..01be11ea 100644
--- a/docs/swagger.json
+++ b/docs/swagger.json
@@ -188,7 +188,8 @@
                             "unbonding_requested",
                             "unbonding",
                             "unbonded",
-                            "withdrawn"
+                            "withdrawn",
+                            "transitioned"
                         ],
                         "type": "string",
                         "description": "Filter by state",
@@ -991,6 +992,9 @@
                 "finality_provider_pk_hex": {
                     "type": "string"
                 },
+                "is_eligible_for_transition": {
+                    "type": "boolean"
+                },
                 "is_overflow": {
                     "type": "boolean"
                 },
diff --git a/docs/swagger.yaml b/docs/swagger.yaml
index c0415519..3c42abbf 100644
--- a/docs/swagger.yaml
+++ b/docs/swagger.yaml
@@ -233,6 +233,8 @@ definitions:
     properties:
       finality_provider_pk_hex:
         type: string
+      is_eligible_for_transition:
+        type: boolean
       is_overflow:
         type: boolean
       staker_pk_hex:
@@ -612,6 +614,7 @@ paths:
         - unbonding
         - unbonded
         - withdrawn
+        - transitioned
         in: query
         name: state
         type: string
diff --git a/internal/indexer/db/client/interface.go b/internal/indexer/db/client/interface.go
index e2c74b0c..bba600e8 100644
--- a/internal/indexer/db/client/interface.go
+++ b/internal/indexer/db/client/interface.go
@@ -21,4 +21,10 @@ type IndexerDBClient interface {
 	// Staker Delegations
 	GetDelegation(ctx context.Context, stakingTxHashHex string) (*indexerdbmodel.IndexerDelegationDetails, error)
 	GetDelegations(ctx context.Context, stakerPKHex string, paginationToken string) (*db.DbResultMap[indexerdbmodel.IndexerDelegationDetails], error)
+	/**
+	 * GetLastProcessedBbnHeight retrieves the last processed BBN height.
+	 * @param ctx The context
+	 * @return The last processed height or an error
+	 */
+	GetLastProcessedBbnHeight(ctx context.Context) (uint64, error)
 }
diff --git a/internal/indexer/db/client/last_processed_height.go b/internal/indexer/db/client/last_processed_height.go
new file mode 100644
index 00000000..0e488020
--- /dev/null
+++ b/internal/indexer/db/client/last_processed_height.go
@@ -0,0 +1,25 @@
+package indexerdbclient
+
+import (
+	"context"
+
+	indexerdbmodel "github.com/babylonlabs-io/staking-api-service/internal/indexer/db/model"
+	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/mongo"
+)
+
+func (db *IndexerDatabase) GetLastProcessedBbnHeight(ctx context.Context) (uint64, error) {
+	// If not in context, query from database
+	var result indexerdbmodel.LastProcessedHeight
+	err := db.Client.Database(db.DbName).Collection(
+		indexerdbmodel.LastProcessedHeightCollection,
+	).FindOne(ctx, bson.M{}).Decode(&result)
+	if err == mongo.ErrNoDocuments {
+		// If no document exists, return 0
+		return 0, nil
+	}
+	if err != nil {
+		return 0, err
+	}
+	return result.Height, nil
+}
diff --git a/internal/indexer/db/model/last_processed_height.go b/internal/indexer/db/model/last_processed_height.go
new file mode 100644
index 00000000..78c17112
--- /dev/null
+++ b/internal/indexer/db/model/last_processed_height.go
@@ -0,0 +1,5 @@
+package indexerdbmodel
+
+type LastProcessedHeight struct {
+	Height uint64 `bson:"height"`
+}
diff --git a/internal/indexer/db/model/setup.go b/internal/indexer/db/model/setup.go
index bf788273..c4e056c4 100644
--- a/internal/indexer/db/model/setup.go
+++ b/internal/indexer/db/model/setup.go
@@ -5,4 +5,5 @@ const (
 	BTCDelegationDetailsCollection    = "btc_delegation_details"
 	TimeLockCollection                = "timelock"
 	GlobalParamsCollection            = "global_params"
+	LastProcessedHeightCollection     = "last_processed_height"
 )
diff --git a/internal/shared/config/config.go b/internal/shared/config/config.go
index cfdef0d7..a80ce513 100644
--- a/internal/shared/config/config.go
+++ b/internal/shared/config/config.go
@@ -10,12 +10,13 @@ import (
 )
 
 type Config struct {
-	Server    *ServerConfig      `mapstructure:"server"`
-	StakingDb *DbConfig          `mapstructure:"staking-db"`
-	IndexerDb *DbConfig          `mapstructure:"indexer-db"`
-	Queue     *queue.QueueConfig `mapstructure:"queue"`
-	Metrics   *MetricsConfig     `mapstructure:"metrics"`
-	Assets    *AssetsConfig      `mapstructure:"assets"`
+	Server               *ServerConfig               `mapstructure:"server"`
+	StakingDb            *DbConfig                   `mapstructure:"staking-db"`
+	IndexerDb            *DbConfig                   `mapstructure:"indexer-db"`
+	Queue                *queue.QueueConfig          `mapstructure:"queue"`
+	Metrics              *MetricsConfig              `mapstructure:"metrics"`
+	Assets               *AssetsConfig               `mapstructure:"assets"`
+	DelegationTransition *DelegationTransitionConfig `mapstructure:"delegation_transition"`
 }
 
 func (cfg *Config) Validate() error {
@@ -46,6 +47,13 @@ func (cfg *Config) Validate() error {
 		}
 	}
 
+	// The DelegationTransition is optional
+	if cfg.DelegationTransition != nil {
+		if err := cfg.DelegationTransition.Validate(); err != nil {
+			return err
+		}
+	}
+
 	return nil
 }
 
diff --git a/internal/shared/config/delegation_transition.go b/internal/shared/config/delegation_transition.go
new file mode 100644
index 00000000..e3702bd4
--- /dev/null
+++ b/internal/shared/config/delegation_transition.go
@@ -0,0 +1,26 @@
+package config
+
+import "errors"
+
+// DelegationTransitionConfig represents the transition cutoff height for
+// the phase-1 delegation to phase-2.
+// A delegation can transition to phase-2 if either:
+// 1. The delegation's BTC staking height is less than EligibleBeforeBtcHeight, or
+// 2. The current BBN height is greater than AllowListExpirationHeight
+// (allowing all delegations to transition)
+type DelegationTransitionConfig struct {
+	EligibleBeforeBtcHeight   uint64 `mapstructure:"eligible_before_btc_height"`
+	AllowListExpirationHeight uint64 `mapstructure:"allow_list_expiration_height"`
+}
+
+func (cfg *DelegationTransitionConfig) Validate() error {
+	if cfg.EligibleBeforeBtcHeight == 0 {
+		return errors.New("before_btc_height cannot be 0")
+	}
+
+	if cfg.AllowListExpirationHeight == 0 {
+		return errors.New("allow_list_expiration_height cannot be 0")
+	}
+
+	return nil
+}
diff --git a/internal/shared/types/delegation.go b/internal/shared/types/delegation.go
index a19e6113..5e687ba7 100644
--- a/internal/shared/types/delegation.go
+++ b/internal/shared/types/delegation.go
@@ -10,6 +10,7 @@ const (
 	Unbonding          DelegationState = "unbonding"
 	Unbonded           DelegationState = "unbonded"
 	Withdrawn          DelegationState = "withdrawn"
+	Transitioned       DelegationState = "transitioned"
 )
 
 func (s DelegationState) ToString() string {
@@ -28,6 +29,8 @@ func FromStringToDelegationState(s string) (DelegationState, error) {
 		return Unbonded, nil
 	case "withdrawn":
 		return Withdrawn, nil
+	case "transitioned":
+		return Transitioned, nil
 	default:
 		return "", fmt.Errorf("invalid delegation state: %s", s)
 	}
diff --git a/internal/v1/api/handlers/delegation.go b/internal/v1/api/handlers/delegation.go
index af23a976..0633086e 100644
--- a/internal/v1/api/handlers/delegation.go
+++ b/internal/v1/api/handlers/delegation.go
@@ -5,7 +5,6 @@ import (
 
 	"github.com/babylonlabs-io/staking-api-service/internal/shared/api/handlers/handler"
 	"github.com/babylonlabs-io/staking-api-service/internal/shared/types"
-	v1service "github.com/babylonlabs-io/staking-api-service/internal/v1/service"
 )
 
 // GetDelegationByTxHash @Summary Get a delegation
@@ -26,5 +25,5 @@ func (h *V1Handler) GetDelegationByTxHash(request *http.Request) (*handler.Resul
 		return nil, err
 	}
 
-	return handler.NewResult(v1service.FromDelegationDocument(delegation)), nil
+	return handler.NewResult(delegation), nil
 }
diff --git a/internal/v1/db/model/delegation.go b/internal/v1/db/model/delegation.go
index c24e3710..df148576 100644
--- a/internal/v1/db/model/delegation.go
+++ b/internal/v1/db/model/delegation.go
@@ -1,8 +1,8 @@
 package v1dbmodel
 
 import (
-	"github.com/babylonlabs-io/staking-api-service/internal/shared/types"
 	dbmodel "github.com/babylonlabs-io/staking-api-service/internal/shared/db/model"
+	"github.com/babylonlabs-io/staking-api-service/internal/shared/types"
 )
 
 type TimelockTransaction struct {
diff --git a/internal/v1/queue/handler/expired_staking.go b/internal/v1/queue/handler/expired_staking.go
index af1fd1ee..f7da1b59 100644
--- a/internal/v1/queue/handler/expired_staking.go
+++ b/internal/v1/queue/handler/expired_staking.go
@@ -26,7 +26,8 @@ func (h *V1QueueHandler) ExpiredStakingHandler(ctx context.Context, messageBody
 	if delErr != nil {
 		return delErr
 	}
-	if utils.Contains[types.DelegationState](utils.OutdatedStatesForUnbonded(), del.State) {
+	state := types.DelegationState(del.State)
+	if utils.Contains(utils.OutdatedStatesForUnbonded(), state) {
 		// Ignore the message as the delegation state already passed the unbonded state. This is an outdated duplication
 		log.Ctx(ctx).Debug().Str("StakingTxHashHex", expiredStakingEvent.StakingTxHashHex).
 			Msg("delegation state is outdated for unbonded event")
diff --git a/internal/v1/queue/handler/unbonding.go b/internal/v1/queue/handler/unbonding.go
index 6e6225f5..7d1ecefb 100644
--- a/internal/v1/queue/handler/unbonding.go
+++ b/internal/v1/queue/handler/unbonding.go
@@ -25,7 +25,7 @@ func (h *V1QueueHandler) UnbondingStakingHandler(ctx context.Context, messageBod
 	if delErr != nil {
 		return delErr
 	}
-	state := del.State
+	state := types.DelegationState(del.State)
 	if utils.Contains(utils.OutdatedStatesForUnbonding(), state) {
 		// Ignore the message as the delegation state already passed the unbonding state. This is an outdated duplication
 		log.Ctx(ctx).Debug().Str("StakingTxHashHex", unbondingStakingEvent.StakingTxHashHex).
diff --git a/internal/v1/queue/handler/withdraw.go b/internal/v1/queue/handler/withdraw.go
index ed15a765..2bd97108 100644
--- a/internal/v1/queue/handler/withdraw.go
+++ b/internal/v1/queue/handler/withdraw.go
@@ -25,7 +25,7 @@ func (h *V1QueueHandler) WithdrawStakingHandler(ctx context.Context, messageBody
 	if delErr != nil {
 		return delErr
 	}
-	state := del.State
+	state := types.DelegationState(del.State)
 
 	stakingTxHashHex := withdrawnStakingEvent.GetStakingTxHashHex()
 
diff --git a/internal/v1/service/delegation.go b/internal/v1/service/delegation.go
index 986cfd0e..cc460ac4 100644
--- a/internal/v1/service/delegation.go
+++ b/internal/v1/service/delegation.go
@@ -21,50 +21,21 @@ type TransactionPublic struct {
 }
 
 type DelegationPublic struct {
-	StakingTxHashHex      string             `json:"staking_tx_hash_hex"`
-	StakerPkHex           string             `json:"staker_pk_hex"`
-	FinalityProviderPkHex string             `json:"finality_provider_pk_hex"`
-	State                 string             `json:"state"`
-	StakingValue          uint64             `json:"staking_value"`
-	StakingTx             *TransactionPublic `json:"staking_tx"`
-	UnbondingTx           *TransactionPublic `json:"unbonding_tx,omitempty"`
-	IsOverflow            bool               `json:"is_overflow"`
-}
-
-func FromDelegationDocument(d *v1model.DelegationDocument) DelegationPublic {
-	delPublic := DelegationPublic{
-		StakingTxHashHex:      d.StakingTxHashHex,
-		StakerPkHex:           d.StakerPkHex,
-		FinalityProviderPkHex: d.FinalityProviderPkHex,
-		StakingValue:          d.StakingValue,
-		State:                 d.State.ToString(),
-		StakingTx: &TransactionPublic{
-			TxHex:          d.StakingTx.TxHex,
-			OutputIndex:    d.StakingTx.OutputIndex,
-			StartTimestamp: utils.ParseTimestampToIsoFormat(d.StakingTx.StartTimestamp),
-			StartHeight:    d.StakingTx.StartHeight,
-			TimeLock:       d.StakingTx.TimeLock,
-		},
-		IsOverflow: d.IsOverflow,
-	}
-
-	// Add unbonding transaction if it exists
-	if d.UnbondingTx != nil && d.UnbondingTx.TxHex != "" {
-		delPublic.UnbondingTx = &TransactionPublic{
-			TxHex:          d.UnbondingTx.TxHex,
-			OutputIndex:    d.UnbondingTx.OutputIndex,
-			StartTimestamp: utils.ParseTimestampToIsoFormat(d.UnbondingTx.StartTimestamp),
-			StartHeight:    d.UnbondingTx.StartHeight,
-			TimeLock:       d.UnbondingTx.TimeLock,
-		}
-	}
-	return delPublic
+	StakingTxHashHex        string             `json:"staking_tx_hash_hex"`
+	StakerPkHex             string             `json:"staker_pk_hex"`
+	FinalityProviderPkHex   string             `json:"finality_provider_pk_hex"`
+	State                   string             `json:"state"`
+	StakingValue            uint64             `json:"staking_value"`
+	StakingTx               *TransactionPublic `json:"staking_tx"`
+	UnbondingTx             *TransactionPublic `json:"unbonding_tx,omitempty"`
+	IsOverflow              bool               `json:"is_overflow"`
+	IsEligibleForTransition bool               `json:"is_eligible_for_transition"`
 }
 
 func (s *V1Service) DelegationsByStakerPk(
 	ctx context.Context, stakerPk string,
 	state types.DelegationState, pageToken string,
-) ([]DelegationPublic, string, *types.Error) {
+) ([]*DelegationPublic, string, *types.Error) {
 	filter := &v1dbclient.DelegationFilter{}
 	if state != "" {
 		filter = &v1dbclient.DelegationFilter{
@@ -81,9 +52,15 @@ func (s *V1Service) DelegationsByStakerPk(
 		log.Ctx(ctx).Error().Err(err).Msg("Failed to find delegations by staker pk")
 		return nil, "", types.NewInternalServiceError(err)
 	}
-	var delegations []DelegationPublic = make([]DelegationPublic, 0, len(resultMap.Data))
+	var delegations []*DelegationPublic = make([]*DelegationPublic, 0, len(resultMap.Data))
+	bbnHeight, err := s.Service.DbClients.IndexerDBClient.GetLastProcessedBbnHeight(ctx)
+	if err != nil {
+		log.Ctx(ctx).Error().Err(err).Msg("Failed to get last processed BBN height")
+		return nil, "", types.NewInternalServiceError(err)
+	}
+
 	for _, d := range resultMap.Data {
-		delegations = append(delegations, FromDelegationDocument(&d))
+		delegations = append(delegations, s.FromDelegationDocument(&d, bbnHeight))
 	}
 	return delegations, resultMap.PaginationToken, nil
 }
@@ -125,7 +102,7 @@ func (s *V1Service) IsDelegationPresent(ctx context.Context, txHashHex string) (
 	return false, nil
 }
 
-func (s *V1Service) GetDelegation(ctx context.Context, txHashHex string) (*v1model.DelegationDocument, *types.Error) {
+func (s *V1Service) GetDelegation(ctx context.Context, txHashHex string) (*DelegationPublic, *types.Error) {
 	delegation, err := s.Service.DbClients.V1DBClient.FindDelegationByTxHashHex(ctx, txHashHex)
 	if err != nil {
 		if db.IsNotFoundError(err) {
@@ -135,7 +112,12 @@ func (s *V1Service) GetDelegation(ctx context.Context, txHashHex string) (*v1mod
 		log.Ctx(ctx).Error().Err(err).Msg("Failed to find delegation by tx hash hex")
 		return nil, types.NewInternalServiceError(err)
 	}
-	return delegation, nil
+	bbnHeight, err := s.Service.DbClients.IndexerDBClient.GetLastProcessedBbnHeight(ctx)
+	if err != nil {
+		log.Ctx(ctx).Error().Err(err).Msg("Failed to get last processed BBN height")
+		return nil, types.NewInternalServiceError(err)
+	}
+	return s.FromDelegationDocument(delegation, bbnHeight), nil
 }
 
 func (s *V1Service) CheckStakerHasActiveDelegationByPk(
@@ -154,3 +136,61 @@ func (s *V1Service) CheckStakerHasActiveDelegationByPk(
 	}
 	return hasDelegation, nil
 }
+
+func (s *V1Service) isEligibleForTransition(
+	delegation *v1model.DelegationDocument, bbnHeight uint64,
+) bool {
+	if s.Cfg.DelegationTransition == nil {
+		return false
+	}
+
+	// Check the delegation state, only active delegations are eligible for transition
+	if delegation.State != types.Active {
+		return false
+	}
+
+	// Check the delegation staking height
+	stakingHeight := delegation.StakingTx.StartHeight
+	// Only not overflow delegations are eligible for transition before the Btc height
+	if !delegation.IsOverflow && stakingHeight < s.Cfg.DelegationTransition.EligibleBeforeBtcHeight {
+		return true
+	}
+	if bbnHeight >= s.Cfg.DelegationTransition.AllowListExpirationHeight {
+		return true
+	}
+
+	return false
+}
+
+func (s *V1Service) FromDelegationDocument(
+	d *v1model.DelegationDocument, bbnHeight uint64,
+) *DelegationPublic {
+	delPublic := &DelegationPublic{
+		StakingTxHashHex:      d.StakingTxHashHex,
+		StakerPkHex:           d.StakerPkHex,
+		FinalityProviderPkHex: d.FinalityProviderPkHex,
+		StakingValue:          d.StakingValue,
+		State:                 d.State.ToString(),
+		StakingTx: &TransactionPublic{
+			TxHex:          d.StakingTx.TxHex,
+			OutputIndex:    d.StakingTx.OutputIndex,
+			StartTimestamp: utils.ParseTimestampToIsoFormat(d.StakingTx.StartTimestamp),
+			StartHeight:    d.StakingTx.StartHeight,
+			TimeLock:       d.StakingTx.TimeLock,
+		},
+		IsOverflow:              d.IsOverflow,
+		IsEligibleForTransition: s.isEligibleForTransition(d, bbnHeight),
+	}
+
+	// Add unbonding transaction if it exists
+	if d.UnbondingTx != nil && d.UnbondingTx.TxHex != "" {
+		delPublic.UnbondingTx = &TransactionPublic{
+			TxHex:          d.UnbondingTx.TxHex,
+			OutputIndex:    d.UnbondingTx.OutputIndex,
+			StartTimestamp: utils.ParseTimestampToIsoFormat(d.UnbondingTx.StartTimestamp),
+			StartHeight:    d.UnbondingTx.StartHeight,
+			TimeLock:       d.UnbondingTx.TimeLock,
+		}
+	}
+	return delPublic
+}
diff --git a/internal/v1/service/interface.go b/internal/v1/service/interface.go
index c80b115c..2206c477 100644
--- a/internal/v1/service/interface.go
+++ b/internal/v1/service/interface.go
@@ -5,16 +5,15 @@ import (
 
 	"github.com/babylonlabs-io/staking-api-service/internal/shared/services/service"
 	"github.com/babylonlabs-io/staking-api-service/internal/shared/types"
-	v1model "github.com/babylonlabs-io/staking-api-service/internal/v1/db/model"
 )
 
 type V1ServiceProvider interface {
 	service.SharedServiceProvider
 	// Delegation
-	DelegationsByStakerPk(ctx context.Context, stakerPk string, state types.DelegationState, pageToken string) ([]DelegationPublic, string, *types.Error)
+	DelegationsByStakerPk(ctx context.Context, stakerPk string, state types.DelegationState, pageToken string) ([]*DelegationPublic, string, *types.Error)
 	SaveActiveStakingDelegation(ctx context.Context, txHashHex, stakerPkHex, finalityProviderPkHex string, value, startHeight uint64, stakingTimestamp int64, timeLock, stakingOutputIndex uint64, stakingTxHex string, isOverflow bool) *types.Error
 	IsDelegationPresent(ctx context.Context, txHashHex string) (bool, *types.Error)
-	GetDelegation(ctx context.Context, txHashHex string) (*v1model.DelegationDocument, *types.Error)
+	GetDelegation(ctx context.Context, txHashHex string) (*DelegationPublic, *types.Error)
 	CheckStakerHasActiveDelegationByPk(ctx context.Context, stakerPkHex string, afterTimestamp int64) (bool, *types.Error)
 	TransitionToUnbondingState(ctx context.Context, txHashHex string, startHeight, timelock, outputIndex uint64, txHex string, startTimestamp int64) *types.Error
 	TransitionToWithdrawnState(ctx context.Context, txHashHex string) *types.Error
diff --git a/tests/mocks/mock_db_client.go b/tests/mocks/mock_db_client.go
index 2911d21c..b247612e 100644
--- a/tests/mocks/mock_db_client.go
+++ b/tests/mocks/mock_db_client.go
@@ -1,4 +1,4 @@
-// Code generated by mockery v2.44.1. DO NOT EDIT.
+// Code generated by mockery v2.41.0. DO NOT EDIT.
 
 package mocks
 
diff --git a/tests/mocks/mock_ordinal_client.go b/tests/mocks/mock_ordinal_client.go
index 1c21cfef..95dff743 100644
--- a/tests/mocks/mock_ordinal_client.go
+++ b/tests/mocks/mock_ordinal_client.go
@@ -1,4 +1,4 @@
-// Code generated by mockery v2.44.1. DO NOT EDIT.
+// Code generated by mockery v2.41.0. DO NOT EDIT.
 
 package mocks
 
diff --git a/tests/mocks/mock_v1_db_client.go b/tests/mocks/mock_v1_db_client.go
index f7492dd9..2d0d206b 100644
--- a/tests/mocks/mock_v1_db_client.go
+++ b/tests/mocks/mock_v1_db_client.go
@@ -1,4 +1,4 @@
-// Code generated by mockery v2.44.1. DO NOT EDIT.
+// Code generated by mockery v2.41.0. DO NOT EDIT.
 
 package mocks
 
diff --git a/tests/mocks/mock_v2_db_client.go b/tests/mocks/mock_v2_db_client.go
index f6c57c26..971f596f 100644
--- a/tests/mocks/mock_v2_db_client.go
+++ b/tests/mocks/mock_v2_db_client.go
@@ -1,4 +1,4 @@
-// Code generated by mockery v2.44.1. DO NOT EDIT.
+// Code generated by mockery v2.41.0. DO NOT EDIT.
 
 package mocks
 
@@ -7,6 +7,8 @@ import (
 
 	dbmodel "github.com/babylonlabs-io/staking-api-service/internal/shared/db/model"
 	mock "github.com/stretchr/testify/mock"
+
+	v2dbmodel "github.com/babylonlabs-io/staking-api-service/internal/v2/db/model"
 )
 
 // V2DBClient is an autogenerated mock type for the V2DBClient type
@@ -122,6 +124,66 @@ func (_m *V2DBClient) FindUnprocessableMessages(ctx context.Context) ([]dbmodel.
 	return r0, r1
 }
 
+// GetOverallStats provides a mock function with given fields: ctx
+func (_m *V2DBClient) GetOverallStats(ctx context.Context) (*v2dbmodel.V2OverallStatsDocument, error) {
+	ret := _m.Called(ctx)
+
+	if len(ret) == 0 {
+		panic("no return value specified for GetOverallStats")
+	}
+
+	var r0 *v2dbmodel.V2OverallStatsDocument
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context) (*v2dbmodel.V2OverallStatsDocument, error)); ok {
+		return rf(ctx)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context) *v2dbmodel.V2OverallStatsDocument); ok {
+		r0 = rf(ctx)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*v2dbmodel.V2OverallStatsDocument)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context) error); ok {
+		r1 = rf(ctx)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// GetStakerStats provides a mock function with given fields: ctx, stakerPKHex
+func (_m *V2DBClient) GetStakerStats(ctx context.Context, stakerPKHex string) (*v2dbmodel.V2StakerStatsDocument, error) {
+	ret := _m.Called(ctx, stakerPKHex)
+
+	if len(ret) == 0 {
+		panic("no return value specified for GetStakerStats")
+	}
+
+	var r0 *v2dbmodel.V2StakerStatsDocument
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, string) (*v2dbmodel.V2StakerStatsDocument, error)); ok {
+		return rf(ctx, stakerPKHex)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, string) *v2dbmodel.V2StakerStatsDocument); ok {
+		r0 = rf(ctx, stakerPKHex)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*v2dbmodel.V2StakerStatsDocument)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, string) error); ok {
+		r1 = rf(ctx, stakerPKHex)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
 // InsertPkAddressMappings provides a mock function with given fields: ctx, stakerPkHex, taproot, nativeSigwitOdd, nativeSigwitEven
 func (_m *V2DBClient) InsertPkAddressMappings(ctx context.Context, stakerPkHex string, taproot string, nativeSigwitOdd string, nativeSigwitEven string) error {
 	ret := _m.Called(ctx, stakerPkHex, taproot, nativeSigwitOdd, nativeSigwitEven)