refactor(x/staking): migrate ValidatorByConsAddr key to collections

tests/integration/auth/migrations/v2/store_test.go

@@ -779,8 +779,11 @@ func createValidator(t *testing.T, ctx sdk.Context, bankKeeper bankkeeper.Keeper
 	val1, err := stakingtypes.NewValidator(valAddrs[0].String(), pks[0], stakingtypes.Description{})
 	require.NoError(t, err)
 
+	consAddr, err := val1.GetConsAddr()
+	require.NoError(t, err)
+
 	require.NoError(t, stakingKeeper.SetValidator(ctx, val1))
-	require.NoError(t, stakingKeeper.SetValidatorByConsAddr(ctx, val1))
+	require.NoError(t, stakingKeeper.ValidatorByConsensusAddress.Set(ctx, consAddr, val1.GetOperator()))
 	require.NoError(t, stakingKeeper.SetNewValidatorByPowerIndex(ctx, val1))
 
 	_, err = stakingKeeper.Delegate(ctx, addrs[0], valTokens, stakingtypes.Unbonded, val1, true)

tests/integration/gov/keeper/common_test.go

@@ -48,12 +48,19 @@ func createValidators(t *testing.T, f *fixture, powers []int64) ([]sdk.AccAddres
 	val3, err := stakingtypes.NewValidator(valAddrs[2].String(), pks[2], stakingtypes.Description{})
 	assert.NilError(t, err)
 
+	consAddr1, err := val1.GetConsAddr()
+	assert.NilError(t, err)
+	consAddr2, err := val2.GetConsAddr()
+	assert.NilError(t, err)
+	consAddr3, err := val3.GetConsAddr()
+	assert.NilError(t, err)
+
 	assert.NilError(t, f.stakingKeeper.SetValidator(f.ctx, val1))
 	assert.NilError(t, f.stakingKeeper.SetValidator(f.ctx, val2))
 	assert.NilError(t, f.stakingKeeper.SetValidator(f.ctx, val3))
-	assert.NilError(t, f.stakingKeeper.SetValidatorByConsAddr(f.ctx, val1))
-	assert.NilError(t, f.stakingKeeper.SetValidatorByConsAddr(f.ctx, val2))
-	assert.NilError(t, f.stakingKeeper.SetValidatorByConsAddr(f.ctx, val3))
+	assert.NilError(t, f.stakingKeeper.ValidatorByConsensusAddress.Set(f.ctx, consAddr1, val1.GetOperator()))
+	assert.NilError(t, f.stakingKeeper.ValidatorByConsensusAddress.Set(f.ctx, consAddr2, val2.GetOperator()))
+	assert.NilError(t, f.stakingKeeper.ValidatorByConsensusAddress.Set(f.ctx, consAddr3, val3.GetOperator()))
 	assert.NilError(t, f.stakingKeeper.SetNewValidatorByPowerIndex(f.ctx, val1))
 	assert.NilError(t, f.stakingKeeper.SetNewValidatorByPowerIndex(f.ctx, val2))
 	assert.NilError(t, f.stakingKeeper.SetNewValidatorByPowerIndex(f.ctx, val3))

tests/integration/staking/keeper/common_test.go

@@ -77,12 +77,20 @@ func createValidators(t *testing.T, f *fixture, powers []int64) ([]sdk.AccAddres
 
 	assert.NilError(t, f.stakingKeeper.SetValidator(f.sdkCtx, val1))
 	assert.NilError(t, f.stakingKeeper.SetValidator(f.sdkCtx, val2))
-	assert.NilError(t, f.stakingKeeper.SetValidatorByConsAddr(f.sdkCtx, val1))
-	assert.NilError(t, f.stakingKeeper.SetValidatorByConsAddr(f.sdkCtx, val2))
+
+	consAddr1, err := val1.GetConsAddr()
+	assert.NilError(t, err)
+
+	consAddr2, err := val2.GetConsAddr()
+	assert.NilError(t, err)
+
+	assert.NilError(t, f.stakingKeeper.ValidatorByConsensusAddress.Set(f.sdkCtx, consAddr1, val1.GetOperator()))
+	assert.NilError(t, f.stakingKeeper.ValidatorByConsensusAddress.Set(f.sdkCtx, consAddr2, val2.GetOperator()))
+
 	assert.NilError(t, f.stakingKeeper.SetNewValidatorByPowerIndex(f.sdkCtx, val1))
 	assert.NilError(t, f.stakingKeeper.SetNewValidatorByPowerIndex(f.sdkCtx, val2))
 
-	_, err := f.stakingKeeper.Delegate(f.sdkCtx, addrs[0], f.stakingKeeper.TokensFromConsensusPower(f.sdkCtx, powers[0]), types.Unbonded, val1, true)
+	_, err = f.stakingKeeper.Delegate(f.sdkCtx, addrs[0], f.stakingKeeper.TokensFromConsensusPower(f.sdkCtx, powers[0]), types.Unbonded, val1, true)
 	assert.NilError(t, err)
 	_, err = f.stakingKeeper.Delegate(f.sdkCtx, addrs[1], f.stakingKeeper.TokensFromConsensusPower(f.sdkCtx, powers[1]), types.Unbonded, val2, true)
 	assert.NilError(t, err)

tests/integration/staking/keeper/deterministic_test.go

@@ -234,16 +234,20 @@ func createAndSetValidator(t *testing.T, rt *rapid.T, f *deterministicFixture) s
 
 func setValidator(t *testing.T, f *deterministicFixture, validator stakingtypes.Validator) {
 	t.Helper()
+
+	consAddr, err := validator.GetConsAddr()
+	assert.NilError(t, err)
+
 	assert.NilError(t, f.stakingKeeper.SetValidator(f.ctx, validator))
 	assert.NilError(t, f.stakingKeeper.SetValidatorByPowerIndex(f.ctx, validator))
-	assert.NilError(t, f.stakingKeeper.SetValidatorByConsAddr(f.ctx, validator))
+	assert.NilError(t, f.stakingKeeper.ValidatorByConsensusAddress.Set(f.ctx, consAddr, validator.GetOperator()))
 	assert.NilError(t, f.stakingKeeper.Hooks().AfterValidatorCreated(f.ctx, validator.GetOperator()))
 
 	delegatorAddress := sdk.AccAddress(validator.GetOperator())
 	coins := sdk.NewCoins(sdk.NewCoin(sdk.DefaultBondDenom, validator.BondedTokens()))
 	assert.NilError(t, banktestutil.FundAccount(f.ctx, f.bankKeeper, delegatorAddress, coins))
 
-	_, err := f.stakingKeeper.Delegate(f.ctx, delegatorAddress, validator.BondedTokens(), stakingtypes.Unbonded, validator, true)
+	_, err = f.stakingKeeper.Delegate(f.ctx, delegatorAddress, validator.BondedTokens(), stakingtypes.Unbonded, validator, true)
 	assert.NilError(t, err)
 }
 

tests/integration/staking/keeper/slash_test.go

@@ -48,7 +48,9 @@ func bootstrapSlashTest(t *testing.T, power int64) (*fixture, []sdk.AccAddress,
 		validator := testutil.NewValidator(t, addrVals[i], PKs[i])
 		validator, _ = validator.AddTokensFromDel(amt)
 		validator = keeper.TestingUpdateValidator(f.stakingKeeper, f.sdkCtx, validator, true)
-		assert.NilError(t, f.stakingKeeper.SetValidatorByConsAddr(f.sdkCtx, validator))
+		consAddr, err := validator.GetConsAddr()
+		assert.NilError(t, err)
+		assert.NilError(t, f.stakingKeeper.ValidatorByConsensusAddress.Set(f.sdkCtx, consAddr, validator.GetOperator()))
 	}
 
 	return f, addrDels, addrVals

tests/integration/staking/keeper/validator_test.go

@@ -144,7 +144,11 @@ func TestSlashToZeroPowerRemoved(t *testing.T) {
 	validator, _ = validator.AddTokensFromDel(valTokens)
 	assert.Equal(t, types.Unbonded, validator.Status)
 	assert.DeepEqual(t, valTokens, validator.Tokens)
-	assert.NilError(t, f.stakingKeeper.SetValidatorByConsAddr(f.sdkCtx, validator))
+
+	consAddr, err := validator.GetConsAddr()
+	assert.NilError(t, err)
+	assert.NilError(t, f.stakingKeeper.ValidatorByConsensusAddress.Set(f.sdkCtx, consAddr, validator.GetOperator()))
+
 	validator = keeper.TestingUpdateValidator(f.stakingKeeper, f.sdkCtx, validator, true)
 	assert.DeepEqual(t, valTokens, validator.Tokens)
 

x/slashing/simulation/operations_test.go

@@ -159,12 +159,13 @@ func (suite *SimTestSuite) TestSimulateMsgUnjail() {
 	validator0, err := getTestingValidator0(ctx, suite.stakingKeeper, suite.accounts)
 	suite.Require().NoError(err)
 
-	// setup validator0 by consensus address
-	err = suite.stakingKeeper.SetValidatorByConsAddr(ctx, validator0)
+	val0ConsAddress, err := validator0.GetConsAddr()
 	suite.Require().NoError(err)
 
-	val0ConsAddress, err := validator0.GetConsAddr()
+	// setup validator0 by consensus address
+	err = suite.stakingKeeper.ValidatorByConsensusAddress.Set(ctx, val0ConsAddress, validator0.GetOperator())
 	suite.Require().NoError(err)
+
 	info := types.NewValidatorSigningInfo(val0ConsAddress, int64(4), int64(3),
 		time.Unix(2, 0), false, int64(10))
 	err = suite.slashingKeeper.ValidatorSigningInfo.Set(ctx, val0ConsAddress, info)

x/staking/keeper/delegation_test.go

@@ -414,7 +414,9 @@ func (s *KeeperTestSuite) TestUndelegateSelfDelegationBelowMinSelfDelegation() {
 
 	s.bankKeeper.EXPECT().SendCoinsFromModuleToModule(gomock.Any(), stakingtypes.NotBondedPoolName, stakingtypes.BondedPoolName, gomock.Any())
 	validator = stakingkeeper.TestingUpdateValidator(keeper, ctx, validator, true)
-	require.NoError(keeper.SetValidatorByConsAddr(ctx, validator))
+	consAddr, err := validator.GetConsAddr()
+	require.NoError(err)
+	require.NoError(keeper.ValidatorByConsensusAddress.Set(ctx, consAddr, validator.GetOperator()))
 	require.True(validator.IsBonded())
 
 	selfDelegation := stakingtypes.NewDelegation(sdk.AccAddress(addrVals[0].Bytes()).String(), addrVals[0].String(), issuedShares)
@@ -432,7 +434,7 @@ func (s *KeeperTestSuite) TestUndelegateSelfDelegationBelowMinSelfDelegation() {
 
 	val0AccAddr := sdk.AccAddress(addrVals[0].Bytes())
 	s.bankKeeper.EXPECT().SendCoinsFromModuleToModule(gomock.Any(), stakingtypes.BondedPoolName, stakingtypes.NotBondedPoolName, gomock.Any())
-	_, _, err := keeper.Undelegate(ctx, val0AccAddr, addrVals[0], math.LegacyNewDecFromInt(keeper.TokensFromConsensusPower(ctx, 6)))
+	_, _, err = keeper.Undelegate(ctx, val0AccAddr, addrVals[0], math.LegacyNewDecFromInt(keeper.TokensFromConsensusPower(ctx, 6)))
 	require.NoError(err)
 
 	// end block
@@ -455,7 +457,9 @@ func (s *KeeperTestSuite) TestUndelegateFromUnbondingValidator() {
 
 	// create a validator with a self-delegation
 	validator := testutil.NewValidator(s.T(), addrVals[0], PKs[0])
-	require.NoError(keeper.SetValidatorByConsAddr(ctx, validator))
+	consAddr, err := validator.GetConsAddr()
+	require.NoError(err)
+	require.NoError(keeper.ValidatorByConsensusAddress.Set(ctx, consAddr, validator.GetOperator()))
 
 	validator, issuedShares := validator.AddTokensFromDel(delTokens)
 	require.Equal(delTokens, issuedShares.RoundInt())
@@ -531,7 +535,9 @@ func (s *KeeperTestSuite) TestUndelegateFromUnbondedValidator() {
 
 	// create a validator with a self-delegation
 	validator := testutil.NewValidator(s.T(), addrVals[0], PKs[0])
-	require.NoError(keeper.SetValidatorByConsAddr(ctx, validator))
+	consAddr, err := validator.GetConsAddr()
+	require.NoError(err)
+	require.NoError(keeper.ValidatorByConsensusAddress.Set(ctx, consAddr, validator.GetOperator()))
 
 	valTokens := keeper.TokensFromConsensusPower(ctx, 10)
 	validator, issuedShares := validator.AddTokensFromDel(valTokens)
@@ -609,7 +615,9 @@ func (s *KeeperTestSuite) TestUnbondingAllDelegationFromValidator() {
 
 	// create a validator with a self-delegation
 	validator := testutil.NewValidator(s.T(), addrVals[0], PKs[0])
-	require.NoError(keeper.SetValidatorByConsAddr(ctx, validator))
+	consAddr, err := validator.GetConsAddr()
+	require.NoError(err)
+	require.NoError(keeper.ValidatorByConsensusAddress.Set(ctx, consAddr, validator.GetOperator()))
 
 	valTokens := keeper.TokensFromConsensusPower(ctx, 10)
 	validator, issuedShares := validator.AddTokensFromDel(valTokens)
@@ -858,7 +866,9 @@ func (s *KeeperTestSuite) TestRedelegateSelfDelegation() {
 
 	// create a validator with a self-delegation
 	validator := testutil.NewValidator(s.T(), addrVals[0], PKs[0])
-	require.NoError(keeper.SetValidatorByConsAddr(ctx, validator))
+	consAddr, err := validator.GetConsAddr()
+	require.NoError(err)
+	require.NoError(keeper.ValidatorByConsensusAddress.Set(ctx, consAddr, validator.GetOperator()))
 
 	valTokens := keeper.TokensFromConsensusPower(ctx, 10)
 	validator, issuedShares := validator.AddTokensFromDel(valTokens)
@@ -888,7 +898,7 @@ func (s *KeeperTestSuite) TestRedelegateSelfDelegation() {
 	delegation := stakingtypes.NewDelegation(addrDels[0].String(), addrVals[0].String(), issuedShares)
 	require.NoError(keeper.SetDelegation(ctx, delegation))
 
-	_, err := keeper.BeginRedelegation(ctx, val0AccAddr, addrVals[0], addrVals[1], math.LegacyNewDecFromInt(delTokens))
+	_, err = keeper.BeginRedelegation(ctx, val0AccAddr, addrVals[0], addrVals[1], math.LegacyNewDecFromInt(delTokens))
 	require.NoError(err)
 
 	// end block
@@ -909,7 +919,9 @@ func (s *KeeperTestSuite) TestRedelegateFromUnbondingValidator() {
 
 	// create a validator with a self-delegation
 	validator := testutil.NewValidator(s.T(), addrVals[0], PKs[0])
-	require.NoError(keeper.SetValidatorByConsAddr(ctx, validator))
+	consAddr, err := validator.GetConsAddr()
+	require.NoError(err)
+	require.NoError(keeper.ValidatorByConsensusAddress.Set(ctx, consAddr, validator.GetOperator()))
 
 	valTokens := keeper.TokensFromConsensusPower(ctx, 10)
 	validator, issuedShares := validator.AddTokensFromDel(valTokens)
@@ -990,7 +1002,9 @@ func (s *KeeperTestSuite) TestRedelegateFromUnbondedValidator() {
 
 	// create a validator with a self-delegation
 	validator := testutil.NewValidator(s.T(), addrVals[0], PKs[0])
-	require.NoError(keeper.SetValidatorByConsAddr(ctx, validator))
+	consAddr, err := validator.GetConsAddr()
+	require.NoError(err)
+	require.NoError(keeper.ValidatorByConsensusAddress.Set(ctx, consAddr, validator.GetOperator()))
 
 	valTokens := keeper.TokensFromConsensusPower(ctx, 10)
 	validator, issuedShares := validator.AddTokensFromDel(valTokens)

x/staking/keeper/genesis.go

@@ -44,7 +44,12 @@ func (k Keeper) InitGenesis(ctx context.Context, data *types.GenesisState) (res
 		}
 
 		// Manually set indices for the first time
-		if err := k.SetValidatorByConsAddr(ctx, validator); err != nil {
+		consAddr, err := validator.GetConsAddr()
+		if err != nil {
+			panic(err)
+		}
+
+		if err := k.ValidatorByConsensusAddress.Set(ctx, consAddr, validator.GetOperator()); err != nil {
 			panic(err)
 		}
 

x/staking/keeper/keeper.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 
 	"cosmossdk.io/collections"
+	collcodec "cosmossdk.io/collections/codec"
 	addresscodec "cosmossdk.io/core/address"
 	storetypes "cosmossdk.io/core/store"
 	"cosmossdk.io/log"
@@ -32,11 +33,12 @@ type Keeper struct {
 	validatorAddressCodec addresscodec.Codec
 	consensusAddressCodec addresscodec.Codec
 
-	Schema                 collections.Schema
-	HistoricalInfo         collections.Map[uint64, types.HistoricalInfo]
-	LastTotalPower         collections.Item[math.Int]
-	ValidatorUpdates       collections.Item[types.ValidatorUpdates]
-	DelegationsByValidator collections.Map[collections.Pair[sdk.ValAddress, sdk.AccAddress], []byte]
+	Schema                      collections.Schema
+	HistoricalInfo              collections.Map[uint64, types.HistoricalInfo]
+	LastTotalPower              collections.Item[math.Int]
+	ValidatorUpdates            collections.Item[types.ValidatorUpdates]
+	DelegationsByValidator      collections.Map[collections.Pair[sdk.ValAddress, sdk.AccAddress], []byte]
+	ValidatorByConsensusAddress collections.Map[sdk.ConsAddress, sdk.ValAddress]
 }
 
 // NewKeeper creates a new staking Keeper instance
@@ -86,6 +88,12 @@ func NewKeeper(
 			collections.PairKeyCodec(sdk.LengthPrefixedAddressKey(sdk.ValAddressKey), sdk.AccAddressKey), // nolint: staticcheck // sdk.LengthPrefixedAddressKey is needed to retain state compatibility
 			collections.BytesValue,
 		),
+		ValidatorByConsensusAddress: collections.NewMap(
+			sb, types.ValidatorsByConsAddrKey,
+			"validator_by_cons_addr",
+			sdk.LengthPrefixedAddressKey(sdk.ConsAddressKey), // nolint: staticcheck // sdk.LengthPrefixedAddressKey is needed to retain state compatibility
+			collcodec.KeyToValueCodec(sdk.ValAddressKey),
+		),
 	}
 
 	schema, err := sb.Build()

x/staking/keeper/msg_server.go

@@ -61,7 +61,8 @@ func (k msgServer) CreateValidator(ctx context.Context, msg *types.MsgCreateVali
 		return nil, errorsmod.Wrapf(sdkerrors.ErrInvalidType, "Expecting cryptotypes.PubKey, got %T", pk)
 	}
 
-	if _, err := k.GetValidatorByConsAddr(ctx, sdk.GetConsAddress(pk)); err == nil {
+	consAddr := sdk.GetConsAddress(pk)
+	if _, err := k.GetValidatorByConsAddr(ctx, consAddr); err == nil {
 		return nil, types.ErrValidatorPubKeyExists
 	}
 
@@ -121,8 +122,7 @@ func (k msgServer) CreateValidator(ctx context.Context, msg *types.MsgCreateVali
 		return nil, err
 	}
 
-	err = k.SetValidatorByConsAddr(ctx, validator)
-	if err != nil {
+	if err := k.ValidatorByConsensusAddress.Set(ctx, consAddr, validator.GetOperator()); err != nil {
 		return nil, err
 	}
 

x/staking/keeper/slash_test.go

@@ -18,7 +18,7 @@ func (s *KeeperTestSuite) TestRevocation() {
 
 	// initial state
 	require.NoError(keeper.SetValidator(ctx, validator))
-	require.NoError(keeper.SetValidatorByConsAddr(ctx, validator))
+	require.NoError(keeper.ValidatorByConsensusAddress.Set(ctx, consAddr, validator.GetOperator()))
 	val, err := keeper.GetValidator(ctx, valAddr)
 	require.NoError(err)
 	require.False(val.IsJailed())
@@ -44,7 +44,7 @@ func (s *KeeperTestSuite) TestSlashAtFutureHeight() {
 	consAddr := sdk.ConsAddress(PKs[0].Address())
 	validator := testutil.NewValidator(s.T(), sdk.ValAddress(PKs[0].Address().Bytes()), PKs[0])
 	require.NoError(keeper.SetValidator(ctx, validator))
-	require.NoError(keeper.SetValidatorByConsAddr(ctx, validator))
+	require.NoError(keeper.ValidatorByConsensusAddress.Set(ctx, consAddr, validator.GetOperator()))
 
 	fraction := sdkmath.LegacyNewDecWithPrec(5, 1)
 	_, err := keeper.Slash(ctx, consAddr, 1, 10, fraction)

x/staking/keeper/validator.go

@@ -9,6 +9,7 @@ import (
 
 	gogotypes "github.com/cosmos/gogoproto/types"
 
+	"cosmossdk.io/collections"
 	corestore "cosmossdk.io/core/store"
 	errorsmod "cosmossdk.io/errors"
 	"cosmossdk.io/math"
@@ -44,9 +45,8 @@ func (k Keeper) mustGetValidator(ctx context.Context, addr sdk.ValAddress) types
 
 // GetValidatorByConsAddr gets a single validator by consensus address
 func (k Keeper) GetValidatorByConsAddr(ctx context.Context, consAddr sdk.ConsAddress) (validator types.Validator, err error) {
-	store := k.storeService.OpenKVStore(ctx)
-	opAddr, err := store.Get(types.GetValidatorByConsAddrKey(consAddr))
-	if err != nil {
+	opAddr, err := k.ValidatorByConsensusAddress.Get(ctx, consAddr)
+	if err != nil && !errors.Is(err, collections.ErrNotFound) {
 		return validator, err
 	}
 
@@ -73,16 +73,6 @@ func (k Keeper) SetValidator(ctx context.Context, validator types.Validator) err
 	return store.Set(types.GetValidatorKey(validator.GetOperator()), bz)
 }
 
-// SetValidatorByConsAddr sets a validator by conesensus address
-func (k Keeper) SetValidatorByConsAddr(ctx context.Context, validator types.Validator) error {
-	consPk, err := validator.GetConsAddr()
-	if err != nil {
-		return err
-	}
-	store := k.storeService.OpenKVStore(ctx)
-	return store.Set(types.GetValidatorByConsAddrKey(consPk), validator.GetOperator())
-}
-
 // SetValidatorByPowerIndex sets a validator by power index
 func (k Keeper) SetValidatorByPowerIndex(ctx context.Context, validator types.Validator) error {
 	// jailed validators are not kept in the power index
@@ -219,7 +209,7 @@ func (k Keeper) RemoveValidator(ctx context.Context, address sdk.ValAddress) err
 		return err
 	}
 
-	if err = store.Delete(types.GetValidatorByConsAddrKey(valConsAddr)); err != nil {
+	if err = k.ValidatorByConsensusAddress.Remove(ctx, valConsAddr); err != nil {
 		return err
 	}
 

x/staking/keeper/validator_test.go

@@ -37,14 +37,17 @@ func (s *KeeperTestSuite) TestValidator() {
 	require.Equal(stakingtypes.Unbonded, validator.Status)
 	require.Equal(valTokens, validator.Tokens)
 	require.Equal(valTokens, validator.DelegatorShares.RoundInt())
+
+	consAddr, err := validator.GetConsAddr()
+	require.NoError(err)
 	require.NoError(keeper.SetValidator(ctx, validator))
 	require.NoError(keeper.SetValidatorByPowerIndex(ctx, validator))
-	require.NoError(keeper.SetValidatorByConsAddr(ctx, validator))
+	require.NoError(keeper.ValidatorByConsensusAddress.Set(ctx, consAddr, validator.GetOperator()))
 
 	// ensure update
 	s.bankKeeper.EXPECT().SendCoinsFromModuleToModule(gomock.Any(), stakingtypes.NotBondedPoolName, stakingtypes.BondedPoolName, gomock.Any())
 	updates := s.applyValidatorSetUpdates(ctx, keeper, 1)
-	validator, err := keeper.GetValidator(ctx, valAddr)
+	validator, err = keeper.GetValidator(ctx, valAddr)
 	require.NoError(err)
 	require.Equal(validator.ABCIValidatorUpdate(keeper.PowerReduction(ctx)), updates[0])
 
@@ -54,8 +57,7 @@ func (s *KeeperTestSuite) TestValidator() {
 	require.Equal(valTokens, validator.DelegatorShares.RoundInt())
 
 	// check each store for being saved
-	consAddr, err := validator.GetConsAddr()
-	require.NoError(err)
+
 	resVal, err := keeper.GetValidatorByConsAddr(ctx, consAddr)
 	require.NoError(err)
 	require.True(validator.MinEqual(&resVal))
@@ -121,7 +123,11 @@ func (s *KeeperTestSuite) TestValidatorBasics() {
 	// set and retrieve a record
 	s.bankKeeper.EXPECT().SendCoinsFromModuleToModule(gomock.Any(), stakingtypes.NotBondedPoolName, stakingtypes.BondedPoolName, gomock.Any())
 	validators[0] = stakingkeeper.TestingUpdateValidator(keeper, ctx, validators[0], true)
-	require.NoError(keeper.SetValidatorByConsAddr(ctx, validators[0]))
+
+	consAddr, err := validators[0].GetConsAddr()
+	require.NoError(err)
+	require.NoError(keeper.ValidatorByConsensusAddress.Set(ctx, consAddr, validators[0].GetOperator()))
+
 	resVal, err := keeper.GetValidator(ctx, sdk.ValAddress(PKs[0].Address().Bytes()))
 	require.NoError(err)
 	require.True(validators[0].MinEqual(&resVal))