From 9c50495dfbbc514440af88f2aebe07cd3e42cee5 Mon Sep 17 00:00:00 2001
From: RafilxTenfen <rafaeltenfen.rt@gmail.com>
Date: Fri, 3 Jan 2025 10:47:18 -0300
Subject: [PATCH] chore: add test that breaks handle rewarding with gaps of
 unfinalized blocks

---
 x/finality/keeper/rewarding.go      |  1 +
 x/finality/keeper/rewarding_test.go | 50 +++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+)

diff --git a/x/finality/keeper/rewarding.go b/x/finality/keeper/rewarding.go
index adcedad9..3b6b3a32 100644
--- a/x/finality/keeper/rewarding.go
+++ b/x/finality/keeper/rewarding.go
@@ -9,6 +9,7 @@ import (
 	"github.com/babylonlabs-io/babylon/x/finality/types"
 )
 
+// HandleRewarding calls the reward to stakers if the block is finalized
 func (k Keeper) HandleRewarding(ctx context.Context, targetHeight int64) {
 	// rewarding is executed in a range of [nextHeightToReward, heightToExamine]
 	// this is we don't know when a block will be finalized and we need ensure
diff --git a/x/finality/keeper/rewarding_test.go b/x/finality/keeper/rewarding_test.go
index 741b4e68..3e644af4 100644
--- a/x/finality/keeper/rewarding_test.go
+++ b/x/finality/keeper/rewarding_test.go
@@ -3,6 +3,7 @@ package keeper_test
 import (
 	"math/rand"
 	"testing"
+	"time"
 
 	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/require"
@@ -108,3 +109,52 @@ func FuzzHandleRewarding(f *testing.F) {
 			"next height should be after second batch of finalized blocks")
 	})
 }
+
+func TestHandleRewardingWithGapsOfUnfinalizedBlocks(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	r := rand.New(rand.NewSource(time.Now().Unix()))
+
+	// Setup keepers
+	bsKeeper := types.NewMockBTCStakingKeeper(ctrl)
+	iKeeper := types.NewMockIncentiveKeeper(ctrl)
+	cKeeper := types.NewMockCheckpointingKeeper(ctrl)
+	fKeeper, ctx := keepertest.FinalityKeeper(t, bsKeeper, iKeeper, cKeeper)
+
+	fpPK, err := datagen.GenRandomBIP340PubKey(r)
+	require.NoError(t, err)
+	fKeeper.SetVotingPower(ctx, fpPK.MustMarshal(), 1, 1)
+
+	// starts rewarding at block 1
+	fKeeper.SetNextHeightToReward(ctx, 1)
+
+	fKeeper.SetBlock(ctx, &types.IndexedBlock{
+		Height:    1,
+		AppHash:   datagen.GenRandomByteArray(r, 32),
+		Finalized: false,
+	})
+	fKeeper.SetBlock(ctx, &types.IndexedBlock{
+		Height:    2,
+		AppHash:   datagen.GenRandomByteArray(r, 32),
+		Finalized: false,
+	})
+
+	// adds the latest finalized block
+	fKeeper.SetBlock(ctx, &types.IndexedBlock{
+		Height:    3,
+		AppHash:   datagen.GenRandomByteArray(r, 32),
+		Finalized: true,
+	})
+	dc := types.NewVotingPowerDistCache()
+	dc.AddFinalityProviderDistInfo(&types.FinalityProviderDistInfo{
+		BtcPk:          fpPK,
+		TotalBondedSat: 1,
+	})
+	fKeeper.SetVotingPowerDistCache(ctx, 3, dc)
+
+	fKeeper.HandleRewarding(ctx, 3)
+
+	actNextBlockToBeRewarded := fKeeper.GetNextHeightToReward(ctx)
+	require.Equal(t, 4, actNextBlockToBeRewarded)
+}