fix: staker delegation pagination not working with txs from same heig…

…ht (#40)

* fix: issue with staker delegation pagination when btc heights are the same

internal/db/delegation.go

@@ -80,7 +80,10 @@ func (db *Database) FindDelegationsByStakerPk(ctx context.Context, stakerPk stri
 	client := db.Client.Database(db.DbName).Collection(model.DelegationCollection)
 
 	filter := bson.M{"staker_pk_hex": stakerPk}
-	options := options.Find().SetSort(bson.M{"staking_tx.start_height": -1}) // Sorting in descending order
+	options := options.Find().SetSort(bson.D{
+		{Key: "staking_tx.start_height", Value: -1},
+		{Key: "_id", Value: 1},
+	})
 
 	// Decode the pagination token first if it exist
 	if paginationToken != "" {

internal/db/model/setup.go

@@ -36,7 +36,7 @@ var collections = map[string][]index{
 	FinalityProviderStatsCollection: {{Indexes: map[string]int{"active_tvl": -1}, Unique: false}},
 	StakerStatsCollection:           {{Indexes: map[string]int{"active_tvl": -1}, Unique: false}},
 	DelegationCollection: {
-		{Indexes: map[string]int{"staker_pk_hex": 1, "staking_tx.start_height": -1}, Unique: false},
+		{Indexes: map[string]int{"staker_pk_hex": 1, "staking_tx.start_height": -1, "_id": 1}, Unique: false},
 		{Indexes: map[string]int{"staker_btc_address.taproot_address": 1, "staking_tx.start_timestamp": -1}, Unique: false},
 	},
 	TimeLockCollection:         {{Indexes: map[string]int{"expire_height": 1}, Unique: false}},

tests/setup.go

@@ -26,12 +26,15 @@ import (
 	"github.com/babylonlabs-io/staking-api-service/internal/clients"
 	"github.com/babylonlabs-io/staking-api-service/internal/config"
 	"github.com/babylonlabs-io/staking-api-service/internal/db"
+	"github.com/babylonlabs-io/staking-api-service/internal/db/model"
 	"github.com/babylonlabs-io/staking-api-service/internal/observability/metrics"
 	"github.com/babylonlabs-io/staking-api-service/internal/queue"
 	"github.com/babylonlabs-io/staking-api-service/internal/services"
 	"github.com/babylonlabs-io/staking-api-service/internal/types"
 )
 
+var setUpDbIndex = false
+
 type TestServerDependency struct {
 	ConfigOverrides         *config.Config
 	MockDbClient            db.DBClient
@@ -169,6 +172,14 @@ func setupTestDB(cfg config.Config) *mongo.Client {
 	if err != nil {
 		log.Fatal(err)
 	}
+	// Setup the db index only once for all tests
+	if !setUpDbIndex {
+		err = model.Setup(context.Background(), &cfg)
+		if err != nil {
+			log.Fatal("Failed to setup database:", err)
+		}
+		setUpDbIndex = true
+	}
 
 	// Purge all collections in the test database
 	if err := PurgeAllCollections(context.TODO(), client, cfg.Db.DbName); err != nil {

tests/staker_test.go

@@ -25,18 +25,25 @@ func FuzzTestStakerDelegationsWithPaginationResponse(f *testing.F) {
 	attachRandomSeedsToFuzzer(f, 3)
 	f.Fuzz(func(t *testing.T, seed int64) {
 		r := rand.New(rand.NewSource(seed))
+		testServer := setupTestServer(t, nil)
+		defer testServer.Close()
+		numOfStaker1Events := int(testServer.Config.Db.MaxPaginationLimit) + r.Intn(100)
 		activeStakingEventsByStaker1 := generateRandomActiveStakingEvents(t, r, &TestActiveEventGeneratorOpts{
-			NumOfEvents:       11,
-			FinalityProviders: generatePks(t, 11),
-			Stakers:           generatePks(t, 1),
+			NumOfEvents: numOfStaker1Events,
+			Stakers:     generatePks(t, 1),
 		})
+		// Different btc height per staking tx
 		activeStakingEventsByStaker2 := generateRandomActiveStakingEvents(t, r, &TestActiveEventGeneratorOpts{
-			NumOfEvents:       11,
-			FinalityProviders: generatePks(t, 11),
-			Stakers:           generatePks(t, 1),
+			NumOfEvents: int(testServer.Config.Db.MaxPaginationLimit) + 1,
+			Stakers:     generatePks(t, 1),
 		})
-		testServer := setupTestServer(t, nil)
-		defer testServer.Close()
+
+		// Modify the height to simulate all events are processed at the same btc height
+		btcHeight := randomBtcHeight(r, 0)
+		for i := range activeStakingEventsByStaker1 {
+			activeStakingEventsByStaker1[i].StakingStartHeight = btcHeight
+		}
+
 		sendTestMessage(
 			testServer.Queues.ActiveStakingQueueClient,
 			append(activeStakingEventsByStaker1, activeStakingEventsByStaker2...),
@@ -45,49 +52,15 @@ func FuzzTestStakerDelegationsWithPaginationResponse(f *testing.F) {
 
 		// Test the API
 		stakerPk := activeStakingEventsByStaker1[0].StakerPkHex
-		url := testServer.Server.URL + stakerDelegations + "?staker_btc_pk=" + stakerPk
-		var paginationKey string
-		var allDataCollected []services.DelegationPublic
-		var atLeastOnePage bool
-		for {
-			resp, err := http.Get(url + "&pagination_key=" + paginationKey)
-			assert.NoError(t, err, "making GET request to delegations by staker pk should not fail")
-			assert.Equal(t, http.StatusOK, resp.StatusCode, "expected HTTP 200 OK status")
-			bodyBytes, err := io.ReadAll(resp.Body)
-			assert.NoError(t, err, "reading response body should not fail")
-			var response handlers.PublicResponse[[]services.DelegationPublic]
-			err = json.Unmarshal(bodyBytes, &response)
-			assert.NoError(t, err, "unmarshalling response body should not fail")
-
-			// Check that the response body is as expected
-			assert.NotEmptyf(t, response.Data, "expected response body to have data")
-			for _, d := range response.Data {
-				assert.Equal(t, stakerPk, d.StakerPkHex, "expected response body to match")
-			}
-			allDataCollected = append(allDataCollected, response.Data...)
-			if response.Pagination.NextKey != "" {
-				paginationKey = response.Pagination.NextKey
-				atLeastOnePage = true
-			} else {
-				break
-			}
-		}
+		fetchPaginatedStakerDelegations(
+			testServer, t, stakerPk, numOfStaker1Events, activeStakingEventsByStaker1,
+		)
 
-		assert.True(t, atLeastOnePage, "expected at least one page of data")
-		assert.Equal(t, 11, len(allDataCollected), "expected 11 items in total")
-		for _, events := range activeStakingEventsByStaker1 {
-			found := false
-			for _, d := range allDataCollected {
-				if d.StakingTxHashHex == events.StakingTxHashHex {
-					found = true
-					break
-				}
-			}
-			assert.True(t, found, "expected to find the staking tx in the response")
-		}
-		for i := 0; i < len(allDataCollected)-1; i++ {
-			assert.True(t, allDataCollected[i].StakingTx.StartHeight >= allDataCollected[i+1].StakingTx.StartHeight, "expected collected data to be sorted by start height")
-		}
+		stakerPk2 := activeStakingEventsByStaker2[0].StakerPkHex
+		fetchPaginatedStakerDelegations(
+			testServer, t, stakerPk2, len(activeStakingEventsByStaker2),
+			activeStakingEventsByStaker2,
+		)
 	})
 }
 
@@ -305,3 +278,50 @@ func fetchCheckStakerActiveDelegations(
 
 	return response.Data
 }
+
+func fetchPaginatedStakerDelegations(
+	testServer *TestServer, t *testing.T, stakerPk string,
+	numOfStakerEvents int, activeStakingEventsByStaker []*client.ActiveStakingEvent,
+) {
+	url := testServer.Server.URL + stakerDelegations + "?staker_btc_pk=" + stakerPk
+	var paginationKey string
+	var allDataCollected []services.DelegationPublic
+	for {
+		resp, err := http.Get(url + "&pagination_key=" + paginationKey)
+		assert.NoError(t, err, "making GET request to delegations by staker pk should not fail")
+		assert.Equal(t, http.StatusOK, resp.StatusCode, "expected HTTP 200 OK status")
+		bodyBytes, err := io.ReadAll(resp.Body)
+		assert.NoError(t, err, "reading response body should not fail")
+		var response handlers.PublicResponse[[]services.DelegationPublic]
+		err = json.Unmarshal(bodyBytes, &response)
+		assert.NoError(t, err, "unmarshalling response body should not fail")
+
+		// Check that the response body is as expected
+		assert.NotEmptyf(t, response.Data, "expected response body to have data")
+		for _, d := range response.Data {
+			assert.Equal(t, stakerPk, d.StakerPkHex, "expected response body to match")
+		}
+		allDataCollected = append(allDataCollected, response.Data...)
+		if response.Pagination.NextKey != "" {
+			paginationKey = response.Pagination.NextKey
+		} else {
+			break
+		}
+	}
+
+	assert.Equal(t, numOfStakerEvents, len(allDataCollected))
+	for _, events := range activeStakingEventsByStaker {
+		found := false
+		for _, d := range allDataCollected {
+			if d.StakingTxHashHex == events.StakingTxHashHex {
+				found = true
+				break
+			}
+		}
+		assert.True(t, found, "expected to find the staking tx in the response")
+	}
+	for i := 0; i < len(allDataCollected)-1; i++ {
+		assert.True(t, allDataCollected[i].StakingTx.StartHeight >=
+			allDataCollected[i+1].StakingTx.StartHeight)
+	}
+}