refactor: make keeper pubkey collection private and add wrappers

x/pubkey/keeper/genesis.go

@@ -3,7 +3,6 @@ package keeper
 import (
 	"bytes"
 
-	"cosmossdk.io/collections"
 	cryptotypes "github.com/cosmos/cosmos-sdk/crypto/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 
@@ -22,7 +21,7 @@ func (k Keeper) InitGenesis(ctx sdk.Context, data types.GenesisState) {
 			if !ok {
 				panic("failed to unpack public key")
 			}
-			err = k.PubKeys.Set(ctx, collections.Join(valAddr, pk.Index), pubKey)
+			err = k.SetValidatorKeyAtIndex(ctx, valAddr, pk.Index, pubKey)
 			if err != nil {
 				panic(err)
 			}
@@ -34,7 +33,7 @@ func (k Keeper) InitGenesis(ctx sdk.Context, data types.GenesisState) {
 func (k Keeper) ExportGenesis(ctx sdk.Context) types.GenesisState {
 	var gs types.GenesisState
 
-	itr, err := k.PubKeys.Iterate(ctx, nil)
+	itr, err := k.pubKeys.Iterate(ctx, nil)
 	if err != nil {
 		panic(err)
 	}

x/pubkey/keeper/keeper.go

@@ -9,6 +9,7 @@ import (
 	"github.com/cosmos/cosmos-sdk/codec"
 	codectypes "github.com/cosmos/cosmos-sdk/codec/types"
 	cryptotypes "github.com/cosmos/cosmos-sdk/crypto/types"
+	sdk "github.com/cosmos/cosmos-sdk/types"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
 
 	"github.com/sedaprotocol/seda-chain/x/pubkey/types"
@@ -19,7 +20,7 @@ type Keeper struct {
 	validatorAddressCodec address.Codec
 
 	Schema  collections.Schema
-	PubKeys collections.Map[collections.Pair[[]byte, uint32], cryptotypes.PubKey]
+	pubKeys collections.Map[collections.Pair[[]byte, uint32], cryptotypes.PubKey]
 }
 
 func NewKeeper(cdc codec.BinaryCodec, storeService storetypes.KVStoreService, sk types.StakingKeeper, validatorAddressCodec address.Codec) *Keeper {
@@ -31,7 +32,7 @@ func NewKeeper(cdc codec.BinaryCodec, storeService storetypes.KVStoreService, sk
 	k := Keeper{
 		stakingKeeper:         sk,
 		validatorAddressCodec: validatorAddressCodec,
-		PubKeys:               collections.NewMap(sb, types.PubKeysPrefix, "pubkeys", collections.PairKeyCodec(collections.BytesKey, collections.Uint32Key), codec.CollInterfaceValue[cryptotypes.PubKey](cdc)),
+		pubKeys:               collections.NewMap(sb, types.PubKeysPrefix, "pubkeys", collections.PairKeyCodec(collections.BytesKey, collections.Uint32Key), codec.CollInterfaceValue[cryptotypes.PubKey](cdc)),
 	}
 
 	schema, err := sb.Build()
@@ -42,14 +43,30 @@ func NewKeeper(cdc codec.BinaryCodec, storeService storetypes.KVStoreService, sk
 	return &k
 }
 
+func (k Keeper) SetValidatorKeyAtIndex(ctx context.Context, validatorAddr sdk.ValAddress, index uint32, pubKey cryptotypes.PubKey) error {
+	err := k.pubKeys.Set(ctx, collections.Join(validatorAddr.Bytes(), index), pubKey)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func (k Keeper) GetValidatorKeyAtIndex(ctx context.Context, validatorAddr sdk.ValAddress, index uint32) (cryptotypes.PubKey, error) {
+	pubKey, err := k.pubKeys.Get(ctx, collections.Join(validatorAddr.Bytes(), index))
+	if err != nil {
+		return nil, err
+	}
+	return pubKey, nil
+}
+
 func (k Keeper) GetValidatorKeys(ctx context.Context, validatorAddr string) (result types.ValidatorPubKeys, err error) {
 	valAddr, err := k.validatorAddressCodec.StringToBytes(validatorAddr)
 	if err != nil {
 		return result, sdkerrors.ErrInvalidAddress.Wrapf("invalid validator address: %s", err)
 	}
 
 	rng := collections.NewPrefixedPairRange[[]byte, uint32](valAddr)
-	itr, err := k.PubKeys.Iterate(ctx, rng)
+	itr, err := k.pubKeys.Iterate(ctx, rng)
 	if err != nil {
 		return result, err
 	}

x/pubkey/keeper/msg_server.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"fmt"
 
-	"cosmossdk.io/collections"
 	cryptotypes "github.com/cosmos/cosmos-sdk/crypto/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
@@ -45,7 +44,7 @@ func (m msgServer) AddKey(goCtx context.Context, msg *types.MsgAddKey) (*types.M
 			return nil, sdkerrors.ErrInvalidType.Wrapf("%T is not a cryptotypes.PubKey", pubKey)
 		}
 
-		if err := m.Keeper.PubKeys.Set(ctx, collections.Join(valAddr, indPubKey.Index), pubKey); err != nil {
+		if err := m.SetValidatorKeyAtIndex(ctx, valAddr, indPubKey.Index, pubKey); err != nil {
 			return nil, err
 		}
 

x/pubkey/keeper/msg_server_test.go

@@ -4,7 +4,6 @@ import (
 	"os"
 	"path/filepath"
 
-	"cosmossdk.io/collections"
 	gomock "go.uber.org/mock/gomock"
 
 	codectypes "github.com/cosmos/cosmos-sdk/codec/types"
@@ -127,7 +126,7 @@ func (s *KeeperTestSuite) TestMsgServer_AddKey() {
 
 			// Check each index.
 			for _, indPubKey := range tt.msg.IndexedPubKeys {
-				pkActual, err := s.keeper.PubKeys.Get(s.ctx, collections.Join(tt.valAddr.Bytes(), indPubKey.Index))
+				pkActual, err := s.keeper.GetValidatorKeyAtIndex(s.ctx, tt.valAddr, indPubKey.Index)
 				s.Require().NoError(err)
 				pkExpected, ok := indPubKey.PubKey.GetCachedValue().(cryptotypes.PubKey)
 				s.Require().True(ok)