chore: use IterateLastValidatorPowers instead of `GetLastValidators…

…` (backport #1953) (#1966) * test: Add integration test reproducing the LastValidators exceeding MaxValidators bug (#1945) * Add test reproducing the LastValidators exceeding MaxValidators * formatting * Update tests/integration/unbonding.go Co-authored-by: insumity <[email protected]> * Update tests/integration/unbonding.go Co-authored-by: Philip Offtermatt <[email protected]> * document --------- Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: insumity <[email protected]> Co-authored-by: Philip Offtermatt <[email protected]> * refactor: use IterateLastValidatorPowers instead of GetLastValidators (#1953) * Add skeleton for GetLastValidators wrapper * Fix unit tests * Correct comment * Log error messages if validators are not found * Change AnyTimes to more specific Times(1) * Instantiate slices with their max length and truncate * Remove GetLastValidators from expectation * Remove GetLastValidators call in consumer * Move GetLastBondedValidators to validator_set_updates * Add comment on iteration loop --------- Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: insumity <[email protected]> Co-authored-by: Philip Offtermatt <[email protected]>
cosmos · Jun 19, 2024 · 134de6a · 134de6a
1 parent 57938b7
commit 134de6a
Show file tree

Hide file tree

Showing 19 changed files with 153 additions and 55 deletions.
diff --git a/tests/integration/unbonding.go b/tests/integration/unbonding.go
@@ -474,9 +474,10 @@ func (s *CCVTestSuite) TestRedelegationProviderFirst() {
 // during the staking module EndBlock.
 func (s *CCVTestSuite) TestTooManyLastValidators() {
 	sk := s.providerApp.GetTestStakingKeeper()
+	pk := s.providerApp.GetProviderKeeper()
 
 	getLastValsFn := func(ctx sdk.Context) []stakingtypes.Validator {
-		lastVals, err := sk.GetLastValidators(s.providerCtx())
+		lastVals, err := pk.GetLastBondedValidators(s.providerCtx())
 		s.Require().NoError(err)
 		return lastVals
 	}

diff --git a/testutil/ibc_testing/generic_setup.go b/testutil/ibc_testing/generic_setup.go
@@ -154,7 +154,7 @@ func AddConsumer[Tp testutil.ProviderApp, Tc testutil.ConsumerApp](
 	s.Require().Len(props, 1, "unexpected len consumer addition proposals in AddConsumer")
 
 	// opt-in all validators
-	lastVals, err := providerApp.GetTestStakingKeeper().GetLastValidators(providerChain.GetContext())
+	lastVals, err := providerApp.GetProviderKeeper().GetLastBondedValidators(providerChain.GetContext())
 	s.Require().NoError(err)
 
 	for _, v := range lastVals {

diff --git a/testutil/keeper/expectations.go b/testutil/keeper/expectations.go
@@ -217,3 +217,37 @@ func GetMocksForSlashValidator(
 			Times(1),
 	}
 }
+
+// SetupMocksForLastBondedValidatorsExpectation sets up the expectation for the `IterateLastValidatorPowers` `MaxValidators`, and `GetValidator` methods of the `mockStakingKeeper` object.
+// These are needed in particular when calling `GetLastBondedValidators` from the provider keeper.
+// Times is the number of times the expectation should be called. Provide -1 for `AnyTimes“.
+func SetupMocksForLastBondedValidatorsExpectation(mockStakingKeeper *MockStakingKeeper, maxValidators uint32, vals []stakingtypes.Validator, powers []int64, times int) {
+	iteratorCall := mockStakingKeeper.EXPECT().IterateLastValidatorPowers(gomock.Any(), gomock.Any()).DoAndReturn(
+		func(ctx sdk.Context, cb func(sdk.ValAddress, int64) bool) error {
+			for i, val := range vals {
+				if stop := cb(sdk.ValAddress(val.OperatorAddress), powers[i]); stop {
+					break
+				}
+			}
+			return nil
+		})
+	maxValidatorsCall := mockStakingKeeper.EXPECT().MaxValidators(gomock.Any()).Return(maxValidators, nil)
+
+	if times == -1 {
+		iteratorCall.AnyTimes()
+		maxValidatorsCall.AnyTimes()
+	} else {
+		iteratorCall.Times(times)
+		maxValidatorsCall.Times(times)
+	}
+
+	// set up mocks for GetValidator calls
+	for _, val := range vals {
+		getValCall := mockStakingKeeper.EXPECT().GetValidator(gomock.Any(), sdk.ValAddress(val.OperatorAddress)).Return(val, nil)
+		if times == -1 {
+			getValCall.AnyTimes()
+		} else {
+			getValCall.Times(times)
+		}
+	}
+}
diff --git a/testutil/keeper/unit_test_helpers.go b/testutil/keeper/unit_test_helpers.go
@@ -227,7 +227,9 @@ func SetupForStoppingConsumerChain(t *testing.T, ctx sdk.Context,
 	providerKeeper *providerkeeper.Keeper, mocks MockedKeepers,
 ) {
 	t.Helper()
-	mocks.MockStakingKeeper.EXPECT().GetLastValidators(gomock.Any()).Times(1)
+
+	SetupMocksForLastBondedValidatorsExpectation(mocks.MockStakingKeeper, 1, []stakingtypes.Validator{}, []int64{}, 1)
+
 	expectations := GetMocksForCreateConsumerClient(ctx, &mocks,
 		"chainID", clienttypes.NewHeight(4, 5))
 	expectations = append(expectations, GetMocksForSetConsumerChain(ctx, &mocks, "chainID")...)

diff --git a/x/ccv/consumer/keeper/changeover_test.go b/x/ccv/consumer/keeper/changeover_test.go
@@ -3,7 +3,6 @@ package keeper_test
 import (
 	"testing"
 
-	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/require"
 
 	sdkcryptocodec "github.com/cosmos/cosmos-sdk/crypto/codec"
@@ -30,6 +29,8 @@ func TestChangeoverToConsumer(t *testing.T) {
 		cIds[4].SDKStakingValidator(),
 	}
 
+	powers := []int64{55, 87324, 2, 42389479, 9089080}
+
 	// Instantiate 5 ics val updates for use in test
 	initialValUpdates := []abci.ValidatorUpdate{
 		{Power: 55, PubKey: cIds[5].TMProtoCryptoPublicKey()},
@@ -41,7 +42,7 @@ func TestChangeoverToConsumer(t *testing.T) {
 
 	testCases := []struct {
 		name string
-		// Last standalone validators that will be mock returned from stakingKeeper.GetLastValidators()
+		// Last standalone validators that will be mock returned from consumerKeeper.GetLastBondedValidators()
 		lastSovVals []stakingtypes.Validator
 		// Val updates corresponding to initial valset set for ccv set initGenesis
 		initialValUpdates []abci.ValidatorUpdate
@@ -100,10 +101,13 @@ func TestChangeoverToConsumer(t *testing.T) {
 		// Set initial valset, as would be done in InitGenesis
 		consumerKeeper.SetInitialValSet(ctx, tc.initialValUpdates)
 
-		// Setup mocked return value for stakingKeeper.GetLastValidators()
-		gomock.InOrder(
-			mocks.MockStakingKeeper.EXPECT().GetLastValidators(ctx).Return(tc.lastSovVals, nil),
-		)
+		// Setup mocked return value for consumerkeeper.GetLastBondedValidators()
+		uthelpers.SetupMocksForLastBondedValidatorsExpectation(
+			mocks.MockStakingKeeper,
+			180, // max validators
+			tc.lastSovVals,
+			powers,
+			-1) // any times
 
 		// Add ref to standalone staking keeper
 		consumerKeeper.SetStandaloneStakingKeeper(mocks.MockStakingKeeper)

diff --git a/x/ccv/consumer/keeper/keeper.go b/x/ccv/consumer/keeper/keeper.go
@@ -353,7 +353,7 @@ func (k Keeper) GetLastStandaloneValidators(ctx sdk.Context) ([]stakingtypes.Val
 	if !k.IsPreCCV(ctx) || k.standaloneStakingKeeper == nil {
 		panic("cannot get last standalone validators if not in pre-ccv state, or if standalone staking keeper is nil")
 	}
-	return k.standaloneStakingKeeper.GetLastValidators(ctx)
+	return k.GetLastBondedValidators(ctx)
 }
 
 // GetElapsedPacketMaturityTimes returns a slice of already elapsed PacketMaturityTimes, sorted by maturity times,
@@ -712,3 +712,9 @@ func (k Keeper) IsPrevStandaloneChain(ctx sdk.Context) bool {
 	store := ctx.KVStore(k.storeKey)
 	return store.Has(types.PrevStandaloneChainKey())
 }
+
+// GetLastBondedValidators iterates the last validator powers in the staking module
+// and returns the first MaxValidators many validators with the largest powers.
+func (k Keeper) GetLastBondedValidators(ctx sdk.Context) ([]stakingtypes.Validator, error) {
+	return ccv.GetLastBondedValidatorsUtil(ctx, k.standaloneStakingKeeper, k.Logger(ctx))
+}
diff --git a/x/ccv/consumer/keeper/keeper_test.go b/x/ccv/consumer/keeper/keeper_test.go
@@ -184,11 +184,14 @@ func TestGetLastSovereignValidators(t *testing.T) {
 	cId1 := crypto.NewCryptoIdentityFromIntSeed(11)
 	val := cId1.SDKStakingValidator()
 	val.Description.Moniker = "sanity check this is the correctly serialized val"
-	gomock.InOrder(
-		mocks.MockStakingKeeper.EXPECT().GetLastValidators(ctx).Return([]stakingtypes.Validator{
-			val,
-		}, nil),
+	testkeeper.SetupMocksForLastBondedValidatorsExpectation(
+		mocks.MockStakingKeeper,
+		180,
+		[]stakingtypes.Validator{val},
+		[]int64{1000},
+		1,
 	)
+
 	lastSovVals, err := ck.GetLastStandaloneValidators(ctx)
 	require.NoError(t, err)
 	require.Equal(t, []stakingtypes.Validator{val}, lastSovVals)

diff --git a/x/ccv/provider/keeper/grpc_query.go b/x/ccv/provider/keeper/grpc_query.go
@@ -10,7 +10,6 @@ import (
 	errorsmod "cosmossdk.io/errors"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
-	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
 
 	"github.com/cosmos/interchain-security/v5/x/ccv/provider/types"
 	ccvtypes "github.com/cosmos/interchain-security/v5/x/ccv/types"
@@ -71,7 +70,7 @@ func (k Keeper) GetConsumerChain(ctx sdk.Context, chainID string) (types.Chain,
 	// Get MinPowerInTop_N
 	var minPowerInTopN int64
 	if found && topN > 0 {
-		bondedValidators, err := k.stakingKeeper.GetLastValidators(ctx)
+		bondedValidators, err := k.GetLastBondedValidators(ctx)
 		if err != nil {
 			return types.Chain{}, err
 		}
@@ -390,9 +389,9 @@ func (k Keeper) hasToValidate(
 	}
 
 	// if the validator was not part of the last epoch, check if the validator is going to be part of te next epoch
-	bondedValidators, err := k.stakingKeeper.GetLastValidators(ctx)
+	bondedValidators, err := k.GetLastBondedValidators(ctx)
 	if err != nil {
-		return false, errorsmod.Wrapf(stakingtypes.ErrNoValidatorFound, "error getting last bonded validators: %s", err)
+		return false, nil
 	}
 	if topN, found := k.GetTopN(ctx, chainID); found && topN > 0 {
 		// in a Top-N chain, we automatically opt in all validators that belong to the top N
@@ -407,7 +406,7 @@ func (k Keeper) hasToValidate(
 	// if the validator is opted in and belongs to the validators of the next epoch, then if nothing changes
 	// the validator would have to validate in the next epoch
 	if k.IsOptedIn(ctx, chainID, provAddr) {
-		lastVals, err := k.stakingKeeper.GetLastValidators(ctx)
+		lastVals, err := k.GetLastBondedValidators(ctx)
 		if err != nil {
 			return false, err
 		}

diff --git a/x/ccv/provider/keeper/grpc_query_test.go b/x/ccv/provider/keeper/grpc_query_test.go
@@ -184,7 +184,7 @@ func TestQueryConsumerChainsValidatorHasToValidate(t *testing.T) {
 	valConsAddr, _ := val.GetConsAddr()
 	providerAddr := types.NewProviderConsAddress(valConsAddr)
 	mocks.MockStakingKeeper.EXPECT().GetValidatorByConsAddr(ctx, valConsAddr).Return(val, nil).AnyTimes()
-	mocks.MockStakingKeeper.EXPECT().GetLastValidators(ctx).Return([]stakingtypes.Validator{val}, nil).AnyTimes()
+	testkeeper.SetupMocksForLastBondedValidatorsExpectation(mocks.MockStakingKeeper, 1, []stakingtypes.Validator{val}, []int64{1}, -1) // -1 to allow the calls "AnyTimes"
 
 	req := types.QueryConsumerChainsValidatorHasToValidateRequest{
 		ProviderAddress: providerAddr.String(),
@@ -203,14 +203,15 @@ func TestQueryConsumerChainsValidatorHasToValidate(t *testing.T) {
 		ConsumerPublicKey: &crypto.PublicKey{
 			Sum: &crypto.PublicKey_Ed25519{
 				Ed25519: []byte{1},
-			}}})
+			},
+		},
+	})
 
 	// set `providerAddr` as an opted-in validator on "chain3"
 	pk.SetOptedIn(ctx, "chain3", providerAddr)
 
 	// `providerAddr` has to validate "chain1" because it is a consumer validator in this chain, as well as "chain3"
-	// because it opted in, in "chain3" and `providerAddr` belongs to the bonded validators (see the mocking of `GetLastValidators`
-	// above)
+	// because it opted in, in "chain3" and `providerAddr` belongs to the bonded validators
 	expectedChains := []string{"chain1", "chain3"}
 
 	res, err := pk.QueryConsumerChainsValidatorHasToValidate(ctx, &req)
@@ -268,7 +269,8 @@ func TestGetConsumerChain(t *testing.T) {
 		{OperatorAddress: "cosmosvaloper1tflk30mq5vgqjdly92kkhhq3raev2hnz6eete3"}, // 500 power
 	}
 	powers := []int64{50, 150, 300, 500} // sum = 1000
-	mocks.MockStakingKeeper.EXPECT().GetLastValidators(gomock.Any()).Return(vals, nil).AnyTimes()
+	maxValidators := uint32(180)
+	testkeeper.SetupMocksForLastBondedValidatorsExpectation(mocks.MockStakingKeeper, maxValidators, vals, powers, -1) // -1 to allow the calls "AnyTimes"
 
 	for i, val := range vals {
 		valAddr, err := sdk.ValAddressFromBech32(val.GetOperator())

diff --git a/x/ccv/provider/keeper/legacy_proposal_test.go b/x/ccv/provider/keeper/legacy_proposal_test.go
@@ -9,6 +9,7 @@ import (
 	"github.com/stretchr/testify/require"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
 
 	testkeeper "github.com/cosmos/interchain-security/v5/testutil/keeper"
 	providerkeeper "github.com/cosmos/interchain-security/v5/x/ccv/provider/keeper"
@@ -106,7 +107,7 @@ func TestHandleLegacyConsumerAdditionProposal(t *testing.T) {
 
 		if tc.expAppendProp {
 			// Mock calls are only asserted if we expect a client to be created.
-			mocks.MockStakingKeeper.EXPECT().GetLastValidators(gomock.Any()).Times(1)
+			testkeeper.SetupMocksForLastBondedValidatorsExpectation(mocks.MockStakingKeeper, 1, []stakingtypes.Validator{}, []int64{}, 1)
 			gomock.InOrder(
 				testkeeper.GetMocksForCreateConsumerClient(ctx, &mocks, tc.prop.ChainId, clienttypes.NewHeight(2, 3))...,
 			)

diff --git a/x/ccv/provider/keeper/partial_set_security.go b/x/ccv/provider/keeper/partial_set_security.go
@@ -69,7 +69,7 @@ func (k Keeper) HandleOptOut(ctx sdk.Context, chainID string, providerAddr types
 		if err != nil {
 			return err
 		}
-		bondedValidators, err := k.stakingKeeper.GetLastValidators(ctx)
+		bondedValidators, err := k.GetLastBondedValidators(ctx)
 		if err != nil {
 			return errorsmod.Wrapf(stakingtypes.ErrNoValidatorFound, "error getting last bonded validators: %s", err)
 		}

diff --git a/x/ccv/provider/keeper/partial_set_security_test.go b/x/ccv/provider/keeper/partial_set_security_test.go
@@ -133,7 +133,7 @@ func TestHandleOptOutFromTopNChain(t *testing.T) {
 	valDConsAddr, _ := valD.GetConsAddr()
 	mocks.MockStakingKeeper.EXPECT().GetValidatorByConsAddr(ctx, valDConsAddr).Return(valD, nil).AnyTimes()
 
-	mocks.MockStakingKeeper.EXPECT().GetLastValidators(ctx).Return([]stakingtypes.Validator{valA, valB, valC, valD}, nil).AnyTimes()
+	testkeeper.SetupMocksForLastBondedValidatorsExpectation(mocks.MockStakingKeeper, 4, []stakingtypes.Validator{valA, valB, valC, valD}, []int64{1, 2, 3, 4}, -1) // -1 to allow mocks AnyTimes
 
 	// opt in all validators
 	providerKeeper.SetOptedIn(ctx, chainID, types.NewProviderConsAddress(valAConsAddr))

diff --git a/x/ccv/provider/keeper/proposal.go b/x/ccv/provider/keeper/proposal.go
@@ -277,7 +277,7 @@ func (k Keeper) MakeConsumerGenesis(
 	}
 
 	// get the bonded validators from the staking module
-	bondedValidators, err := k.stakingKeeper.GetLastValidators(ctx)
+	bondedValidators, err := k.GetLastBondedValidators(ctx)
 	if err != nil {
 		return gen, nil, errorsmod.Wrapf(stakingtypes.ErrNoValidatorFound, "error getting last bonded validators: %s", err)
 	}

diff --git a/x/ccv/provider/keeper/proposal_test.go b/x/ccv/provider/keeper/proposal_test.go
@@ -118,7 +118,7 @@ func TestHandleConsumerAdditionProposal(t *testing.T) {
 
 		if tc.expAppendProp {
 			// Mock calls are only asserted if we expect a client to be created.
-			mocks.MockStakingKeeper.EXPECT().GetLastValidators(gomock.Any()).Times(1)
+			testkeeper.SetupMocksForLastBondedValidatorsExpectation(mocks.MockStakingKeeper, 0, []stakingtypes.Validator{}, []int64{}, 1) // returns empty validator set
 			gomock.InOrder(
 				testkeeper.GetMocksForCreateConsumerClient(ctx, &mocks, tc.prop.ChainId, clienttypes.NewHeight(2, 3))...,
 			)
@@ -162,7 +162,7 @@ func TestCreateConsumerClient(t *testing.T) {
 			description: "No state mutation, new client should be created",
 			setup: func(providerKeeper *providerkeeper.Keeper, ctx sdk.Context, mocks *testkeeper.MockedKeepers) {
 				// Valid client creation is asserted with mock expectations here
-				mocks.MockStakingKeeper.EXPECT().GetLastValidators(gomock.Any()).Times(1)
+				testkeeper.SetupMocksForLastBondedValidatorsExpectation(mocks.MockStakingKeeper, 0, []stakingtypes.Validator{}, []int64{}, 1) // returns empty validator set
 				gomock.InOrder(
 					testkeeper.GetMocksForCreateConsumerClient(ctx, mocks, "chainID", clienttypes.NewHeight(4, 5))...,
 				)
@@ -178,7 +178,7 @@ func TestCreateConsumerClient(t *testing.T) {
 				mocks.MockStakingKeeper.EXPECT().UnbondingTime(gomock.Any()).Times(0)
 				mocks.MockClientKeeper.EXPECT().CreateClient(gomock.Any(), gomock.Any(), gomock.Any()).Times(0)
 				mocks.MockClientKeeper.EXPECT().GetSelfConsensusState(gomock.Any(), gomock.Any()).Times(0)
-				mocks.MockStakingKeeper.EXPECT().GetLastValidators(gomock.Any()).Times(0)
+				testkeeper.SetupMocksForLastBondedValidatorsExpectation(mocks.MockStakingKeeper, 0, []stakingtypes.Validator{}, []int64{}, 0) // returns empty validator set
 			},
 			expClientCreated: false,
 		},
@@ -664,7 +664,7 @@ func TestMakeConsumerGenesis(t *testing.T) {
 	//
 	ctx = ctx.WithChainID("testchain1") // chainID is obtained from ctx
 	ctx = ctx.WithBlockHeight(5)        // RevisionHeight obtained from ctx
-	mocks.MockStakingKeeper.EXPECT().GetLastValidators(gomock.Any()).Times(1)
+	testkeeper.SetupMocksForLastBondedValidatorsExpectation(mocks.MockStakingKeeper, 0, []stakingtypes.Validator{}, []int64{}, 1)
 	gomock.InOrder(testkeeper.GetMocksForMakeConsumerGenesis(ctx, &mocks, 1814400000000000)...)
 
 	// matches params from jsonString
@@ -910,9 +910,10 @@ func TestBeginBlockInit(t *testing.T) {
 	// opt in a sample validator so the chain's proposal can successfully execute
 	validator := cryptotestutil.NewCryptoIdentityFromIntSeed(0).SDKStakingValidator()
 	consAddr, _ := validator.GetConsAddr()
-	mocks.MockStakingKeeper.EXPECT().GetLastValidators(gomock.Any()).Return([]stakingtypes.Validator{validator}, nil).AnyTimes()
-	valAdrr, _ := sdk.ValAddressFromBech32(validator.GetOperator())
-	mocks.MockStakingKeeper.EXPECT().GetLastValidatorPower(gomock.Any(), valAdrr).Return(int64(1), nil).AnyTimes()
+	testkeeper.SetupMocksForLastBondedValidatorsExpectation(mocks.MockStakingKeeper, 1, []stakingtypes.Validator{validator}, []int64{0}, -1) // -1 to allow any number of calls
+
+	valAddr, _ := sdk.ValAddressFromBech32(validator.GetOperator())
+	mocks.MockStakingKeeper.EXPECT().GetLastValidatorPower(gomock.Any(), valAddr).Return(int64(1), nil).AnyTimes()
 	providerKeeper.SetOptedIn(ctx, pendingProps[4].ChainId, providertypes.NewProviderConsAddress(consAddr))
 
 	providerKeeper.BeginBlockInit(ctx)
@@ -1017,7 +1018,7 @@ func TestBeginBlockCCR(t *testing.T) {
 	expectations := []*gomock.Call{}
 	for _, prop := range pendingProps {
 		// A consumer chain is setup corresponding to each prop, making these mocks necessary
-		mocks.MockStakingKeeper.EXPECT().GetLastValidators(gomock.Any()).Times(1)
+		testkeeper.SetupMocksForLastBondedValidatorsExpectation(mocks.MockStakingKeeper, 0, []stakingtypes.Validator{}, []int64{}, 1)
 		expectations = append(expectations, testkeeper.GetMocksForCreateConsumerClient(ctx, &mocks,
 			prop.ChainId, clienttypes.NewHeight(2, 3))...)
 		expectations = append(expectations, testkeeper.GetMocksForSetConsumerChain(ctx, &mocks, prop.ChainId)...)

diff --git a/x/ccv/provider/keeper/relay.go b/x/ccv/provider/keeper/relay.go
@@ -216,11 +216,13 @@ func (k Keeper) SendVSCPacketsToChain(ctx sdk.Context, chainID, channelID string
 
 // QueueVSCPackets queues latest validator updates for every registered consumer chain
 // failing to GetLastValidators will cause a panic in EndBlock
+
+// TODO: decide if this func shouldn't return an error to be propagated to BeginBlocker
 func (k Keeper) QueueVSCPackets(ctx sdk.Context) {
 	valUpdateID := k.GetValidatorSetUpdateId(ctx) // current valset update ID
 
 	// get the bonded validators from the staking module
-	bondedValidators, err := k.stakingKeeper.GetLastValidators(ctx)
+	bondedValidators, err := k.GetLastBondedValidators(ctx)
 	if err != nil {
 		panic(fmt.Errorf("failed to get last validators: %w", err))
 	}

diff --git a/x/ccv/provider/keeper/relay_test.go b/x/ccv/provider/keeper/relay_test.go
@@ -68,18 +68,7 @@ func TestQueueVSCPackets(t *testing.T) {
 		ctrl := gomock.NewController(t)
 		defer ctrl.Finish()
 		mocks := testkeeper.NewMockedKeepers(ctrl)
-		// TODO: remove this after main merge
-		// mockStakingKeeper := mocks.MockStakingKeeper
-
-		// mockUpdates := []abci.ValidatorUpdate{}
-		// if len(tc.packets) != 0 {
-		// 	mockUpdates = tc.packets[0].ValidatorUpdates
-		// }
-
-		// gomock.InOrder(
-		// 	mockStakingKeeper.EXPECT().GetValidatorUpdates(gomock.Eq(ctx)).Return(mockUpdates, nil),
-		// )
-		mocks.MockStakingKeeper.EXPECT().GetLastValidators(ctx).Times(1)
+		testkeeper.SetupMocksForLastBondedValidatorsExpectation(mocks.MockStakingKeeper, 0, []stakingtypes.Validator{}, []int64{}, 1)
 
 		pk := testkeeper.NewInMemProviderKeeper(keeperParams, mocks)
 		// no-op if tc.packets is empty
@@ -772,15 +761,17 @@ func TestEndBlockVSU(t *testing.T) {
 
 	// create 4 sample lastValidators
 	var lastValidators []stakingtypes.Validator
+	var powers []int64
 	for i := 0; i < 4; i++ {
 		validator := crypto.NewCryptoIdentityFromIntSeed(i).SDKStakingValidator()
 		lastValidators = append(lastValidators, validator)
 		valAdrr, err := sdk.ValAddressFromBech32(validator.GetOperator())
 		require.NoError(t, err)
 		mocks.MockStakingKeeper.EXPECT().GetLastValidatorPower(gomock.Any(), valAdrr).Return(int64(i+1), nil).AnyTimes()
+		powers = append(powers, int64(i+1))
 	}
 
-	mocks.MockStakingKeeper.EXPECT().GetLastValidators(gomock.Any()).Return(lastValidators, nil).AnyTimes()
+	testkeeper.SetupMocksForLastBondedValidatorsExpectation(mocks.MockStakingKeeper, 5, lastValidators, powers, -1)
 
 	// set a sample client for a consumer chain so that `GetAllConsumerChains` in `QueueVSCPackets` iterates at least once
 	providerKeeper.SetConsumerClientId(ctx, chainID, "clientID")
@@ -835,7 +826,7 @@ func TestQueueVSCPacketsWithPowerCapping(t *testing.T) {
 	valEPubKey, _ := valE.TmConsPublicKey()
 	mocks.MockStakingKeeper.EXPECT().GetValidatorByConsAddr(ctx, valEConsAddr).Return(valE, nil).AnyTimes()
 
-	mocks.MockStakingKeeper.EXPECT().GetLastValidators(ctx).Return([]stakingtypes.Validator{valA, valB, valC, valD, valE}, nil).AnyTimes()
+	testkeeper.SetupMocksForLastBondedValidatorsExpectation(mocks.MockStakingKeeper, 5, []stakingtypes.Validator{valA, valB, valC, valD, valE}, []int64{1, 3, 4, 8, 16}, -1)
 
 	// add a consumer chain
 	providerKeeper.SetConsumerClientId(ctx, "chainID", "clientID")