test(integration): port types tests to server v2

tests/integration/types/filtered_pagination_test.go → tests/integration/v2/types/filtered_pagination_test.go

@@ -1,17 +1,16 @@
-package query_test
+package types
 
 import (
+	"context"
 	"fmt"
 
+	"cosmossdk.io/collections"
 	"cosmossdk.io/math"
-	"cosmossdk.io/store/prefix"
-	storetypes "cosmossdk.io/store/types"
+	bankkeeper "cosmossdk.io/x/bank/keeper"
 	"cosmossdk.io/x/bank/testutil"
 	"cosmossdk.io/x/bank/types"
 
-	"github.com/cosmos/cosmos-sdk/codec"
 	sdk "github.com/cosmos/cosmos-sdk/types"
-	"github.com/cosmos/cosmos-sdk/types/address"
 	"github.com/cosmos/cosmos-sdk/types/query"
 )
 
@@ -34,17 +33,16 @@ func (s *paginationTestSuite) TestFilteredPaginations() {
 	acc1 := s.accountKeeper.NewAccountWithAddress(s.ctx, addr1)
 	s.accountKeeper.SetAccount(s.ctx, acc1)
 	s.Require().NoError(testutil.FundAccount(s.ctx, s.bankKeeper, addr1, balances))
-	store := s.ctx.KVStore(s.app.UnsafeFindStoreKey(types.StoreKey))
 
 	// verify pagination with limit > total values
 	pageReq := &query.PageRequest{Key: nil, Limit: 5, CountTotal: true}
-	balances, res, err := execFilterPaginate(store, pageReq, s.cdc)
+	balances, res, err := execFilterPaginate(s.ctx, s.bankKeeper, pageReq)
 	s.Require().NoError(err)
 	s.Require().NotNil(res)
 	s.Require().Equal(4, len(balances))
 
 	s.T().Log("verify empty request")
-	balances, res, err = execFilterPaginate(store, nil, s.cdc)
+	balances, res, err = execFilterPaginate(s.ctx, s.bankKeeper, pageReq)
 	s.Require().NoError(err)
 	s.Require().NotNil(res)
 	s.Require().Equal(4, len(balances))
@@ -53,7 +51,7 @@ func (s *paginationTestSuite) TestFilteredPaginations() {
 
 	s.T().Log("verify nextKey is returned if there are more results")
 	pageReq = &query.PageRequest{Key: nil, Limit: 2, CountTotal: true}
-	balances, res, err = execFilterPaginate(store, pageReq, s.cdc)
+	balances, res, err = execFilterPaginate(s.ctx, s.bankKeeper, pageReq)
 	s.Require().NoError(err)
 	s.Require().NotNil(res)
 	s.Require().Equal(2, len(balances))
@@ -63,28 +61,28 @@ func (s *paginationTestSuite) TestFilteredPaginations() {
 
 	s.T().Log("verify both key and offset can't be given")
 	pageReq = &query.PageRequest{Key: res.NextKey, Limit: 1, Offset: 2, CountTotal: true}
-	_, _, err = execFilterPaginate(store, pageReq, s.cdc)
+	_, _, err = execFilterPaginate(s.ctx, s.bankKeeper, pageReq)
 	s.Require().Error(err)
 
 	s.T().Log("use nextKey for query")
 	pageReq = &query.PageRequest{Key: res.NextKey, Limit: 2, CountTotal: true}
-	balances, res, err = execFilterPaginate(store, pageReq, s.cdc)
+	balances, res, err = execFilterPaginate(s.ctx, s.bankKeeper, pageReq)
 	s.Require().NoError(err)
 	s.Require().NotNil(res)
 	s.Require().Equal(2, len(balances))
 	s.Require().Nil(res.NextKey)
 
 	s.T().Log("verify default limit")
 	pageReq = &query.PageRequest{Key: nil, Limit: 0}
-	balances, res, err = execFilterPaginate(store, pageReq, s.cdc)
+	balances, res, err = execFilterPaginate(s.ctx, s.bankKeeper, pageReq)
 	s.Require().NoError(err)
 	s.Require().NotNil(res)
 	s.Require().Equal(4, len(balances))
 	s.Require().Equal(uint64(4), res.Total)
 
 	s.T().Log("verify with offset")
 	pageReq = &query.PageRequest{Offset: 2, Limit: 2}
-	balances, res, err = execFilterPaginate(store, pageReq, s.cdc)
+	balances, res, err = execFilterPaginate(s.ctx, s.bankKeeper, pageReq)
 	s.Require().NoError(err)
 	s.Require().NotNil(res)
 	s.Require().LessOrEqual(len(balances), 2)
@@ -107,17 +105,16 @@ func (s *paginationTestSuite) TestReverseFilteredPaginations() {
 	acc1 := s.accountKeeper.NewAccountWithAddress(s.ctx, addr1)
 	s.accountKeeper.SetAccount(s.ctx, acc1)
 	s.Require().NoError(testutil.FundAccount(s.ctx, s.bankKeeper, addr1, balances))
-	store := s.ctx.KVStore(s.app.UnsafeFindStoreKey(types.StoreKey))
 
 	// verify pagination with limit > total values
 	pageReq := &query.PageRequest{Key: nil, Limit: 5, CountTotal: true, Reverse: true}
-	balns, res, err := execFilterPaginate(store, pageReq, s.cdc)
+	balns, res, err := execFilterPaginate(s.ctx, s.bankKeeper, pageReq)
 	s.Require().NoError(err)
 	s.Require().NotNil(res)
 	s.Require().Equal(5, len(balns))
 
 	s.T().Log("verify empty request")
-	balns, res, err = execFilterPaginate(store, nil, s.cdc)
+	balns, res, err = execFilterPaginate(s.ctx, s.bankKeeper, nil)
 	s.Require().NoError(err)
 	s.Require().NotNil(res)
 	s.Require().Equal(10, len(balns))
@@ -126,15 +123,15 @@ func (s *paginationTestSuite) TestReverseFilteredPaginations() {
 
 	s.T().Log("verify default limit")
 	pageReq = &query.PageRequest{Reverse: true}
-	balns, res, err = execFilterPaginate(store, pageReq, s.cdc)
+	balns, res, err = execFilterPaginate(s.ctx, s.bankKeeper, pageReq)
 	s.Require().NoError(err)
 	s.Require().NotNil(res)
 	s.Require().Equal(10, len(balns))
 	s.Require().Equal(uint64(10), res.Total)
 
 	s.T().Log("verify nextKey is returned if there are more results")
 	pageReq = &query.PageRequest{Limit: 2, CountTotal: true, Reverse: true}
-	balns, res, err = execFilterPaginate(store, pageReq, s.cdc)
+	balns, res, err = execFilterPaginate(s.ctx, s.bankKeeper, pageReq)
 	s.Require().NoError(err)
 	s.Require().NotNil(res)
 	s.Require().Equal(2, len(balns))
@@ -144,12 +141,12 @@ func (s *paginationTestSuite) TestReverseFilteredPaginations() {
 
 	s.T().Log("verify both key and offset can't be given")
 	pageReq = &query.PageRequest{Key: res.NextKey, Limit: 1, Offset: 2, Reverse: true}
-	_, _, err = execFilterPaginate(store, pageReq, s.cdc)
+	_, _, err = execFilterPaginate(s.ctx, s.bankKeeper, pageReq)
 	s.Require().Error(err)
 
 	s.T().Log("use nextKey for query and reverse true")
 	pageReq = &query.PageRequest{Key: res.NextKey, Limit: 2, Reverse: true}
-	balns, res, err = execFilterPaginate(store, pageReq, s.cdc)
+	balns, res, err = execFilterPaginate(s.ctx, s.bankKeeper, pageReq)
 	s.Require().NoError(err)
 	s.Require().NotNil(res)
 	s.Require().Equal(2, len(balns))
@@ -158,7 +155,7 @@ func (s *paginationTestSuite) TestReverseFilteredPaginations() {
 
 	s.T().Log("verify last page records, nextKey for query and reverse true")
 	pageReq = &query.PageRequest{Key: res.NextKey, Reverse: true}
-	balns, res, err = execFilterPaginate(store, pageReq, s.cdc)
+	balns, res, err = execFilterPaginate(s.ctx, s.bankKeeper, pageReq)
 	s.Require().NoError(err)
 	s.Require().NotNil(res)
 	s.Require().Equal(6, len(balns))
@@ -190,62 +187,23 @@ func (s *paginationTestSuite) TestFilteredPaginate() {
 	}
 
 	pageReq := &query.PageRequest{Key: nil, Limit: 1, CountTotal: true}
-	store := s.ctx.KVStore(s.app.UnsafeFindStoreKey(types.StoreKey))
-	balancesStore := prefix.NewStore(store, types.BalancesPrefix)
-	accountStore := prefix.NewStore(balancesStore, address.MustLengthPrefix(addr1))
 
 	var balResult sdk.Coins
-	pageRes, err := query.FilteredPaginate(accountStore, pageReq, func(key, value []byte, accumulate bool) (bool, error) {
-		var amount math.Int
-		err := amount.Unmarshal(value)
-		if err != nil {
-			return false, err
-		}
-
-		// filter amount greater than 100
-		if amount.Int64() > int64(100) {
-			if accumulate {
-				balResult = append(balResult, sdk.NewCoin(string(key), amount))
-			}
-
-			return true, nil
-		}
-
-		return false, nil
-	})
-	if err != nil { // should return no error
-		fmt.Println(err)
-	}
+	balResult, pageRes, err := execFilterPaginate(s.ctx, s.bankKeeper, pageReq)
+	s.Require().NoError(err)
 	fmt.Println(&types.QueryAllBalancesResponse{Balances: balResult, Pagination: pageRes})
 	// Output:
 	// balances:<denom:"test0denom" amount:"250" > pagination:<next_key:"test1denom" total:5 >
 }
 
-func execFilterPaginate(store storetypes.KVStore, pageReq *query.PageRequest, _ codec.Codec) (balances sdk.Coins, res *query.PageResponse, err error) {
-	balancesStore := prefix.NewStore(store, types.BalancesPrefix)
-	accountStore := prefix.NewStore(balancesStore, address.MustLengthPrefix(addr1))
-
-	var balResult sdk.Coins
-	res, err = query.FilteredPaginate(accountStore, pageReq, func(key, value []byte, accumulate bool) (bool, error) {
-		var amount math.Int
-		err := amount.Unmarshal(value)
-		if err != nil {
-			return false, err
-		}
-
+func execFilterPaginate(ctx context.Context, bk bankkeeper.BaseKeeper, pageReq *query.PageRequest) (balances sdk.Coins, res *query.PageResponse, err error) {
+	return query.CollectionFilteredPaginate(ctx, bk.Balances, pageReq, func(key collections.Pair[sdk.AccAddress, string], amount math.Int) (bool, error) {
 		// filter amount greater than 100
-		if amount.Int64() > int64(100) {
-			if accumulate {
-				balResult = append(balResult, sdk.NewCoin(string(key), amount))
-			}
-
-			return true, nil
-		}
-
-		return false, nil
-	})
-
-	return balResult, res, err
+		return amount.Int64() > int64(100), nil
+	}, func(key collections.Pair[sdk.AccAddress, string], amount math.Int) (sdk.Coin, error) {
+		balance := sdk.NewCoin(key.K2(), amount)
+		return balance, nil
+	}, query.WithCollectionPaginationPairPrefix[sdk.AccAddress, string](addr1))
 }
 
 func (s *paginationTestSuite) TestFilteredPaginationsNextKey() {
@@ -261,38 +219,18 @@ func (s *paginationTestSuite) TestFilteredPaginationsNextKey() {
 	acc1 := s.accountKeeper.NewAccountWithAddress(s.ctx, addr1)
 	s.accountKeeper.SetAccount(s.ctx, acc1)
 	s.Require().NoError(testutil.FundAccount(s.ctx, s.bankKeeper, addr1, balances))
-	store := s.ctx.KVStore(s.app.UnsafeFindStoreKey(types.StoreKey))
-
-	execFilterPaginate := func(store storetypes.KVStore, pageReq *query.PageRequest, _ codec.Codec) (balances sdk.Coins, res *query.PageResponse, err error) {
-		balancesStore := prefix.NewStore(store, types.BalancesPrefix)
-		accountStore := prefix.NewStore(balancesStore, address.MustLengthPrefix(addr1))
-
-		var balResult sdk.Coins
-		res, err = query.FilteredPaginate(accountStore, pageReq, func(key, value []byte, accumulate bool) (bool, error) {
-			var amount math.Int
-			err := amount.Unmarshal(value)
-			if err != nil {
-				return false, err
-			}
-
-			// filter odd amounts
-			if amount.Int64()%2 == 1 {
-				if accumulate {
-					balResult = append(balResult, sdk.NewCoin(string(key), amount))
-				}
-
-				return true, nil
-			}
-
-			return false, nil
-		})
-
-		return balResult, res, err
+	execFilterPaginate := func(ctx context.Context, bk bankkeeper.BaseKeeper, pageReq *query.PageRequest) (balances sdk.Coins, res *query.PageResponse, err error) {
+		return query.CollectionFilteredPaginate(ctx, bk.Balances, pageReq, func(key collections.Pair[sdk.AccAddress, string], amount math.Int) (bool, error) {
+			return amount.Int64()%2 == 1, nil
+		}, func(key collections.Pair[sdk.AccAddress, string], amount math.Int) (sdk.Coin, error) {
+			balance := sdk.NewCoin(key.K2(), amount)
+			return balance, nil
+		}, query.WithCollectionPaginationPairPrefix[sdk.AccAddress, string](addr1))
 	}
 
 	s.T().Log("verify next key of offset query")
 	pageReq := &query.PageRequest{Key: nil, Limit: 1, CountTotal: true}
-	balances, res, err := execFilterPaginate(store, pageReq, s.cdc)
+	balances, res, err := execFilterPaginate(s.ctx, s.bankKeeper, pageReq)
 	s.Require().NoError(err)
 	s.Require().NotNil(res)
 	s.Require().Equal(1, len(balances))
@@ -301,7 +239,7 @@ func (s *paginationTestSuite) TestFilteredPaginationsNextKey() {
 	s.Require().NotNil(res.NextKey)
 
 	pageReq = &query.PageRequest{Key: res.NextKey, Limit: 1}
-	balances, res, err = execFilterPaginate(store, pageReq, s.cdc)
+	balances, res, err = execFilterPaginate(s.ctx, s.bankKeeper, pageReq)
 	s.Require().NoError(err)
 	s.Require().NotNil(res)
 	s.Require().Equal(1, len(balances))

tests/integration/types/fuzz_test.go → tests/integration/v2/types/fuzz_test.go

@@ -1,18 +1,16 @@
-package query_test
+package types
 
 import (
 	"fmt"
 	"testing"
 
 	fuzz "github.com/google/gofuzz"
 
+	"cosmossdk.io/collections"
 	"cosmossdk.io/math"
-	"cosmossdk.io/store/prefix"
 	"cosmossdk.io/x/bank/testutil"
 
-	sdktestutil "github.com/cosmos/cosmos-sdk/testutil"
 	sdk "github.com/cosmos/cosmos-sdk/types"
-	"github.com/cosmos/cosmos-sdk/types/address"
 	"github.com/cosmos/cosmos-sdk/types/query"
 )
 
@@ -84,18 +82,9 @@ func FuzzPagination(f *testing.F) {
 		}
 
 		// Now try to paginate it.
-		balResult := sdk.NewCoins()
-		authStore := suite.ctx.KVStore(suite.app.UnsafeFindStoreKey(sdktestutil.BankModuleName))
-		balancesStore := prefix.NewStore(authStore, []byte{balancesPrefix})
-		accountStore := prefix.NewStore(balancesStore, address.MustLengthPrefix(addr1))
-		_, _ = query.Paginate(accountStore, qr, func(key, value []byte) error {
-			var amount math.Int
-			err := amount.Unmarshal(value)
-			if err != nil {
-				return err
-			}
-			balResult = append(balResult, sdk.NewCoin(string(key), amount))
-			return nil
-		})
+		_, _, _ = query.CollectionPaginate(suite.ctx, suite.bankKeeper.Balances, qr, func(key collections.Pair[sdk.AccAddress, string], amount math.Int) (sdk.Coin, error) {
+			balance := sdk.NewCoin(key.K2(), amount)
+			return balance, nil
+		}, query.WithCollectionPaginationPairPrefix[sdk.AccAddress, string](addr1))
 	})
 }

tests/integration/types/pagination_test.go → tests/integration/v2/types/pagination_test.go

@@ -1,36 +1,33 @@
-package query_test
+package types
 
 import (
-	gocontext "context"
+	"context"
 	"fmt"
 	"testing"
 
-	cmtproto "github.com/cometbft/cometbft/api/cometbft/types/v1"
 	"github.com/stretchr/testify/suite"
 
+	"cosmossdk.io/collections"
 	"cosmossdk.io/depinject"
 	"cosmossdk.io/log"
 	"cosmossdk.io/math"
-	"cosmossdk.io/store/prefix"
 	_ "cosmossdk.io/x/accounts"
 	_ "cosmossdk.io/x/bank"
 	bankkeeper "cosmossdk.io/x/bank/keeper"
 	"cosmossdk.io/x/bank/testutil"
 	"cosmossdk.io/x/bank/types"
 	_ "cosmossdk.io/x/consensus"
 
-	"github.com/cosmos/cosmos-sdk/baseapp"
 	"github.com/cosmos/cosmos-sdk/codec"
 	codectypes "github.com/cosmos/cosmos-sdk/codec/types"
 	"github.com/cosmos/cosmos-sdk/crypto/keys/secp256k1"
-	"github.com/cosmos/cosmos-sdk/runtime"
+	"github.com/cosmos/cosmos-sdk/tests/integration/v2"
 	"github.com/cosmos/cosmos-sdk/testutil/configurator"
-	testutilsims "github.com/cosmos/cosmos-sdk/testutil/sims"
 	sdk "github.com/cosmos/cosmos-sdk/types"
-	"github.com/cosmos/cosmos-sdk/types/address"
 	"github.com/cosmos/cosmos-sdk/types/query"
 	_ "github.com/cosmos/cosmos-sdk/x/auth"
 	authkeeper "github.com/cosmos/cosmos-sdk/x/auth/keeper"
+	_ "github.com/cosmos/cosmos-sdk/x/auth/tx/config"
 )
 
 const (
@@ -46,44 +43,50 @@ const (
 
 type paginationTestSuite struct {
 	suite.Suite
-	ctx           sdk.Context
-	bankKeeper    bankkeeper.Keeper
+	ctx           context.Context
+	bankKeeper    bankkeeper.BaseKeeper
 	accountKeeper authkeeper.AccountKeeper
 	cdc           codec.Codec
 	interfaceReg  codectypes.InterfaceRegistry
-	app           *runtime.App
+	app           *integration.App
+
+	queryClient bankkeeper.Querier
 }
 
 func TestPaginationTestSuite(t *testing.T) {
 	suite.Run(t, new(paginationTestSuite))
 }
 
 func (s *paginationTestSuite) SetupTest() {
-	var (
-		bankKeeper    bankkeeper.Keeper
-		accountKeeper authkeeper.AccountKeeper
-		reg           codectypes.InterfaceRegistry
-		cdc           codec.Codec
-	)
+	moduleConfigs := []configurator.ModuleOption{
+		configurator.AccountsModule(),
+		configurator.AuthModule(),
+		configurator.BankModule(),
+		configurator.TxModule(),
+		configurator.ValidateModule(),
+		configurator.ConsensusModule(),
+		configurator.OmitInitGenesis(),
+	}
 
-	app, err := testutilsims.Setup(
+	var err error
+	startupCfg := integration.DefaultStartUpConfig(s.T())
+
+	startupCfg.BranchService = &integration.BranchService{}
+	startupCfg.HeaderService = &integration.HeaderService{}
+
+	s.app, err = integration.NewApp(
 		depinject.Configs(
-			configurator.NewAppConfig(
-				configurator.AccountsModule(),
-				configurator.AuthModule(),
-				configurator.BankModule(),
-				configurator.ConsensusModule(),
-				configurator.OmitInitGenesis(),
-			),
+			configurator.NewAppV2Config(moduleConfigs...),
 			depinject.Supply(log.NewNopLogger()),
 		),
-		&bankKeeper, &accountKeeper, &reg, &cdc)
-
-	s.NoError(err)
+		startupCfg,
+		&s.bankKeeper, &s.accountKeeper, &s.cdc, &s.interfaceReg,
+	)
+	s.Assert().NoError(err)
 
-	ctx := app.BaseApp.NewContextLegacy(false, cmtproto.Header{Height: 1})
+	s.ctx = s.app.StateLatestContext(s.T())
 
-	s.ctx, s.bankKeeper, s.accountKeeper, s.cdc, s.app, s.interfaceReg = ctx, bankKeeper, accountKeeper, cdc, app, reg
+	s.queryClient = bankkeeper.NewQuerier(&s.bankKeeper)
 
 	s.Require().NoError(s.bankKeeper.SetParams(s.ctx, types.DefaultParams()))
 }
@@ -108,9 +111,7 @@ func (s *paginationTestSuite) TestParsePagination() {
 }
 
 func (s *paginationTestSuite) TestPagination() {
-	queryHelper := baseapp.NewQueryServerTestHelper(s.ctx, s.interfaceReg)
-	types.RegisterQueryServer(queryHelper, s.bankKeeper)
-	queryClient := types.NewQueryClient(queryHelper)
+	queryClient := s.queryClient
 
 	var balances sdk.Coins
 
@@ -131,7 +132,7 @@ func (s *paginationTestSuite) TestPagination() {
 	s.T().Log("verify empty page request results a max of defaultLimit records and counts total records")
 	pageReq := &query.PageRequest{}
 	request := types.NewQueryAllBalancesRequest(addr1Str, pageReq, false)
-	res, err := queryClient.AllBalances(gocontext.Background(), request)
+	res, err := queryClient.AllBalances(s.ctx, request)
 	s.Require().NoError(err)
 	s.Require().Equal(res.Pagination.Total, uint64(numBalances))
 	s.Require().NotNil(res.Pagination.NextKey)
@@ -140,7 +141,7 @@ func (s *paginationTestSuite) TestPagination() {
 	s.T().Log("verify page request with limit > defaultLimit, returns less or equal to `limit` records")
 	pageReq = &query.PageRequest{Limit: overLimit}
 	request = types.NewQueryAllBalancesRequest(addr1Str, pageReq, false)
-	res, err = queryClient.AllBalances(gocontext.Background(), request)
+	res, err = queryClient.AllBalances(s.ctx, request)
 	s.Require().NoError(err)
 	s.Require().Equal(res.Pagination.Total, uint64(0))
 	s.Require().NotNil(res.Pagination.NextKey)
@@ -149,7 +150,7 @@ func (s *paginationTestSuite) TestPagination() {
 	s.T().Log("verify paginate with custom limit and countTotal true")
 	pageReq = &query.PageRequest{Limit: underLimit, CountTotal: true}
 	request = types.NewQueryAllBalancesRequest(addr1Str, pageReq, false)
-	res, err = queryClient.AllBalances(gocontext.Background(), request)
+	res, err = queryClient.AllBalances(s.ctx, request)
 	s.Require().NoError(err)
 	s.Require().Equal(res.Balances.Len(), underLimit)
 	s.Require().NotNil(res.Pagination.NextKey)
@@ -158,7 +159,7 @@ func (s *paginationTestSuite) TestPagination() {
 	s.T().Log("verify paginate with custom limit and countTotal false")
 	pageReq = &query.PageRequest{Limit: defaultLimit, CountTotal: false}
 	request = types.NewQueryAllBalancesRequest(addr1Str, pageReq, false)
-	res, err = queryClient.AllBalances(gocontext.Background(), request)
+	res, err = queryClient.AllBalances(s.ctx, request)
 	s.Require().NoError(err)
 	s.Require().Equal(res.Balances.Len(), defaultLimit)
 	s.Require().NotNil(res.Pagination.NextKey)
@@ -167,7 +168,7 @@ func (s *paginationTestSuite) TestPagination() {
 	s.T().Log("verify paginate with custom limit, key and countTotal false")
 	pageReq = &query.PageRequest{Key: res.Pagination.NextKey, Limit: defaultLimit, CountTotal: false}
 	request = types.NewQueryAllBalancesRequest(addr1Str, pageReq, false)
-	res, err = queryClient.AllBalances(gocontext.Background(), request)
+	res, err = queryClient.AllBalances(s.ctx, request)
 	s.Require().NoError(err)
 	s.Require().Equal(res.Balances.Len(), defaultLimit)
 	s.Require().NotNil(res.Pagination.NextKey)
@@ -176,7 +177,7 @@ func (s *paginationTestSuite) TestPagination() {
 	s.T().Log("verify paginate for last page, results in records less than max limit")
 	pageReq = &query.PageRequest{Key: res.Pagination.NextKey, Limit: defaultLimit, CountTotal: false}
 	request = types.NewQueryAllBalancesRequest(addr1Str, pageReq, false)
-	res, err = queryClient.AllBalances(gocontext.Background(), request)
+	res, err = queryClient.AllBalances(s.ctx, request)
 	s.Require().NoError(err)
 	s.Require().LessOrEqual(res.Balances.Len(), defaultLimit)
 	s.Require().Equal(res.Balances.Len(), lastPageRecords)
@@ -186,7 +187,7 @@ func (s *paginationTestSuite) TestPagination() {
 	s.T().Log("verify paginate with offset and limit")
 	pageReq = &query.PageRequest{Offset: 200, Limit: defaultLimit, CountTotal: false}
 	request = types.NewQueryAllBalancesRequest(addr1Str, pageReq, false)
-	res, err = queryClient.AllBalances(gocontext.Background(), request)
+	res, err = queryClient.AllBalances(s.ctx, request)
 	s.Require().NoError(err)
 	s.Require().LessOrEqual(res.Balances.Len(), defaultLimit)
 	s.Require().Equal(res.Balances.Len(), lastPageRecords)
@@ -196,7 +197,7 @@ func (s *paginationTestSuite) TestPagination() {
 	s.T().Log("verify paginate with offset and limit")
 	pageReq = &query.PageRequest{Offset: 100, Limit: defaultLimit, CountTotal: false}
 	request = types.NewQueryAllBalancesRequest(addr1Str, pageReq, false)
-	res, err = queryClient.AllBalances(gocontext.Background(), request)
+	res, err = queryClient.AllBalances(s.ctx, request)
 	s.Require().NoError(err)
 	s.Require().LessOrEqual(res.Balances.Len(), defaultLimit)
 	s.Require().NotNil(res.Pagination.NextKey)
@@ -205,23 +206,21 @@ func (s *paginationTestSuite) TestPagination() {
 	s.T().Log("verify paginate with offset and key - error")
 	pageReq = &query.PageRequest{Key: res.Pagination.NextKey, Offset: 100, Limit: defaultLimit, CountTotal: false}
 	request = types.NewQueryAllBalancesRequest(addr1Str, pageReq, false)
-	_, err = queryClient.AllBalances(gocontext.Background(), request)
+	_, err = queryClient.AllBalances(s.ctx, request)
 	s.Require().Error(err)
 	s.Require().Equal("rpc error: code = InvalidArgument desc = paginate: invalid request, either offset or key is expected, got both", err.Error())
 
 	s.T().Log("verify paginate with offset greater than total results")
 	pageReq = &query.PageRequest{Offset: 300, Limit: defaultLimit, CountTotal: false}
 	request = types.NewQueryAllBalancesRequest(addr1Str, pageReq, false)
-	res, err = queryClient.AllBalances(gocontext.Background(), request)
+	res, err = queryClient.AllBalances(s.ctx, request)
 	s.Require().NoError(err)
 	s.Require().LessOrEqual(res.Balances.Len(), 0)
 	s.Require().Nil(res.Pagination.NextKey)
 }
 
 func (s *paginationTestSuite) TestReversePagination() {
-	queryHelper := baseapp.NewQueryServerTestHelper(s.ctx, s.interfaceReg)
-	types.RegisterQueryServer(queryHelper, s.bankKeeper)
-	queryClient := types.NewQueryClient(queryHelper)
+	queryClient := s.queryClient
 
 	var balances sdk.Coins
 
@@ -242,15 +241,15 @@ func (s *paginationTestSuite) TestReversePagination() {
 	s.T().Log("verify paginate with custom limit and countTotal, Reverse false")
 	pageReq := &query.PageRequest{Limit: 2, CountTotal: true, Reverse: true, Key: nil}
 	request := types.NewQueryAllBalancesRequest(addr1Str, pageReq, false)
-	res1, err := queryClient.AllBalances(gocontext.Background(), request)
+	res1, err := queryClient.AllBalances(s.ctx, request)
 	s.Require().NoError(err)
 	s.Require().Equal(2, res1.Balances.Len())
 	s.Require().NotNil(res1.Pagination.NextKey)
 
 	s.T().Log("verify paginate with custom limit and countTotal, Reverse false")
 	pageReq = &query.PageRequest{Limit: 150}
 	request = types.NewQueryAllBalancesRequest(addr1Str, pageReq, false)
-	res1, err = queryClient.AllBalances(gocontext.Background(), request)
+	res1, err = queryClient.AllBalances(s.ctx, request)
 	s.Require().NoError(err)
 	s.Require().Equal(res1.Balances.Len(), 150)
 	s.Require().NotNil(res1.Pagination.NextKey)
@@ -259,7 +258,7 @@ func (s *paginationTestSuite) TestReversePagination() {
 	s.T().Log("verify paginate with custom limit, key and Reverse true")
 	pageReq = &query.PageRequest{Limit: defaultLimit, Reverse: true}
 	request = types.NewQueryAllBalancesRequest(addr1Str, pageReq, false)
-	res, err := queryClient.AllBalances(gocontext.Background(), request)
+	res, err := queryClient.AllBalances(s.ctx, request)
 	s.Require().NoError(err)
 	s.Require().Equal(res.Balances.Len(), defaultLimit)
 	s.Require().NotNil(res.Pagination.NextKey)
@@ -268,7 +267,7 @@ func (s *paginationTestSuite) TestReversePagination() {
 	s.T().Log("verify paginate with custom limit, key and Reverse true")
 	pageReq = &query.PageRequest{Offset: 100, Limit: defaultLimit, Reverse: true}
 	request = types.NewQueryAllBalancesRequest(addr1Str, pageReq, false)
-	res, err = queryClient.AllBalances(gocontext.Background(), request)
+	res, err = queryClient.AllBalances(s.ctx, request)
 	s.Require().NoError(err)
 	s.Require().Equal(res.Balances.Len(), defaultLimit)
 	s.Require().NotNil(res.Pagination.NextKey)
@@ -277,7 +276,7 @@ func (s *paginationTestSuite) TestReversePagination() {
 	s.T().Log("verify paginate for last page, Reverse true")
 	pageReq = &query.PageRequest{Offset: 200, Limit: defaultLimit, Reverse: true}
 	request = types.NewQueryAllBalancesRequest(addr1Str, pageReq, false)
-	res, err = queryClient.AllBalances(gocontext.Background(), request)
+	res, err = queryClient.AllBalances(s.ctx, request)
 	s.Require().NoError(err)
 	s.Require().Equal(res.Balances.Len(), lastPageRecords)
 	s.Require().Nil(res.Pagination.NextKey)
@@ -286,7 +285,7 @@ func (s *paginationTestSuite) TestReversePagination() {
 	s.T().Log("verify page request with limit > defaultLimit, returns less or equal to `limit` records")
 	pageReq = &query.PageRequest{Limit: overLimit, Reverse: true}
 	request = types.NewQueryAllBalancesRequest(addr1Str, pageReq, false)
-	res, err = queryClient.AllBalances(gocontext.Background(), request)
+	res, err = queryClient.AllBalances(s.ctx, request)
 	s.Require().NoError(err)
 	s.Require().Equal(res.Pagination.Total, uint64(0))
 	s.Require().NotNil(res.Pagination.NextKey)
@@ -295,7 +294,7 @@ func (s *paginationTestSuite) TestReversePagination() {
 	s.T().Log("verify paginate with custom limit, key, countTotal false and Reverse true")
 	pageReq = &query.PageRequest{Key: res1.Pagination.NextKey, Limit: 50, Reverse: true}
 	request = types.NewQueryAllBalancesRequest(addr1Str, pageReq, false)
-	res, err = queryClient.AllBalances(gocontext.Background(), request)
+	res, err = queryClient.AllBalances(s.ctx, request)
 	s.Require().NoError(err)
 	s.Require().Equal(res.Balances.Len(), 50)
 	s.Require().NotNil(res.Pagination.NextKey)
@@ -307,7 +306,7 @@ func (s *paginationTestSuite) TestReversePagination() {
 	s.T().Log("verify paginate with custom limit, key, countTotal false and Reverse true")
 	pageReq = &query.PageRequest{Key: res.Pagination.NextKey, Limit: 50, Reverse: true}
 	request = types.NewQueryAllBalancesRequest(addr1Str, pageReq, false)
-	res, err = queryClient.AllBalances(gocontext.Background(), request)
+	res, err = queryClient.AllBalances(s.ctx, request)
 	s.Require().NoError(err)
 	s.Require().Equal(res.Balances.Len(), 50)
 	s.Require().NotNil(res.Pagination.NextKey)
@@ -319,7 +318,7 @@ func (s *paginationTestSuite) TestReversePagination() {
 	s.T().Log("verify paginate for last page Reverse true")
 	pageReq = &query.PageRequest{Key: res.Pagination.NextKey, Limit: defaultLimit, Reverse: true}
 	request = types.NewQueryAllBalancesRequest(addr1Str, pageReq, false)
-	res, err = queryClient.AllBalances(gocontext.Background(), request)
+	res, err = queryClient.AllBalances(s.ctx, request)
 	s.Require().NoError(err)
 	s.Require().Equal(res.Balances.Len(), 51)
 	s.Require().Nil(res.Pagination.NextKey)
@@ -331,14 +330,14 @@ func (s *paginationTestSuite) TestReversePagination() {
 	s.T().Log("verify paginate with offset and key - error")
 	pageReq = &query.PageRequest{Key: res1.Pagination.NextKey, Offset: 100, Limit: defaultLimit, CountTotal: false}
 	request = types.NewQueryAllBalancesRequest(addr1Str, pageReq, false)
-	_, err = queryClient.AllBalances(gocontext.Background(), request)
+	_, err = queryClient.AllBalances(s.ctx, request)
 	s.Require().Error(err)
 	s.Require().Equal("rpc error: code = InvalidArgument desc = paginate: invalid request, either offset or key is expected, got both", err.Error())
 
 	s.T().Log("verify paginate with offset greater than total results")
 	pageReq = &query.PageRequest{Offset: 300, Limit: defaultLimit, CountTotal: false, Reverse: true}
 	request = types.NewQueryAllBalancesRequest(addr1Str, pageReq, false)
-	res, err = queryClient.AllBalances(gocontext.Background(), request)
+	res, err = queryClient.AllBalances(s.ctx, request)
 	s.Require().NoError(err)
 	s.Require().LessOrEqual(res.Balances.Len(), 0)
 	s.Require().Nil(res.Pagination.NextKey)
@@ -367,22 +366,11 @@ func (s *paginationTestSuite) TestPaginate() {
 	// Paginate example
 	pageReq := &query.PageRequest{Key: nil, Limit: 1, CountTotal: true}
 	request := types.NewQueryAllBalancesRequest(addr1Str, pageReq, false)
-	balResult := sdk.NewCoins()
-	authStore := s.ctx.KVStore(s.app.UnsafeFindStoreKey(types.StoreKey))
-	balancesStore := prefix.NewStore(authStore, types.BalancesPrefix)
-	accountStore := prefix.NewStore(balancesStore, address.MustLengthPrefix(addr1))
-	pageRes, err := query.Paginate(accountStore, request.Pagination, func(key, value []byte) error {
-		var amount math.Int
-		err := amount.Unmarshal(value)
-		if err != nil {
-			return err
-		}
-		balResult = append(balResult, sdk.NewCoin(string(key), amount))
-		return nil
-	})
-	if err != nil { // should return no error
-		fmt.Println(err)
-	}
+	balResult, pageRes, err := query.CollectionPaginate(s.ctx, s.bankKeeper.Balances, request.Pagination, func(key collections.Pair[sdk.AccAddress, string], amount math.Int) (sdk.Coin, error) {
+		balance := sdk.NewCoin(key.K2(), amount)
+		return balance, nil
+	}, query.WithCollectionPaginationPairPrefix[sdk.AccAddress, string](addr1))
+	s.Require().NoError(err)
 	fmt.Println(&types.QueryAllBalancesResponse{Balances: balResult, Pagination: pageRes})
 	// Output:
 	// balances:<denom:"foo0denom" amount:"100" > pagination:<next_key:"foo1denom" total:2 >