From ed00e1a87b1c635c61a4a6c7449c5956ca99ef13 Mon Sep 17 00:00:00 2001
From: zsystm <124245155+zsystm@users.noreply.github.com>
Date: Fri, 20 Sep 2024 10:47:24 +0900
Subject: [PATCH] fix(evmstaking): query withdrawal queue (#113)

it should append QueueElementsPrefixSuffix when query elements of queue
it fixes https://github.com/piplabs/story/issues/84
---
 client/x/evmstaking/keeper/grpc_query.go      |  3 +-
 client/x/evmstaking/keeper/grpc_query_test.go | 63 +++++++++++++++++++
 client/x/evmstaking/keeper/keeper_test.go     |  5 ++
 3 files changed, 70 insertions(+), 1 deletion(-)
 create mode 100644 client/x/evmstaking/keeper/grpc_query_test.go

diff --git a/client/x/evmstaking/keeper/grpc_query.go b/client/x/evmstaking/keeper/grpc_query.go
index 13c942bd..7b1b13f2 100644
--- a/client/x/evmstaking/keeper/grpc_query.go
+++ b/client/x/evmstaking/keeper/grpc_query.go
@@ -8,6 +8,7 @@ import (
 	"github.com/cosmos/cosmos-sdk/runtime"
 	"github.com/cosmos/cosmos-sdk/types/query"
 
+	"github.com/piplabs/story/client/collections"
 	"github.com/piplabs/story/client/x/evmstaking/types"
 
 	"google.golang.org/grpc/codes"
@@ -37,7 +38,7 @@ func (k Keeper) GetWithdrawalQueue(ctx context.Context, request *types.QueryGetW
 	}
 
 	store := runtime.KVStoreAdapter(k.storeService.OpenKVStore(ctx))
-	wqStore := prefix.NewStore(store, types.WithdrawalQueueKey) // withdrawal queue store
+	wqStore := prefix.NewStore(store, append(types.WithdrawalQueueKey, collections.QueueElementsPrefixSuffix)) // withdrawal queue store
 
 	withdrawals, pageResp, err := query.GenericFilteredPaginate(k.cdc, wqStore, request.Pagination, func(_ []byte, wit *types.Withdrawal) (*types.Withdrawal, error) {
 		return wit, nil
diff --git a/client/x/evmstaking/keeper/grpc_query_test.go b/client/x/evmstaking/keeper/grpc_query_test.go
new file mode 100644
index 00000000..25d15a12
--- /dev/null
+++ b/client/x/evmstaking/keeper/grpc_query_test.go
@@ -0,0 +1,63 @@
+package keeper_test
+
+import (
+	"context"
+
+	"github.com/cosmos/cosmos-sdk/types/query"
+	"github.com/ethereum/go-ethereum/common"
+
+	"github.com/piplabs/story/client/x/evmstaking/types"
+)
+
+func (s *TestSuite) TestGetWithdrawalQueue() {
+	require := s.Require()
+	ctx, keeper, queryClient := s.Ctx, s.EVMStakingKeeper, s.queryClient
+	require.NoError(keeper.WithdrawalQueue.Initialize(ctx))
+
+	pageReq := &query.PageRequest{
+		Key:        nil,
+		Limit:      2,
+		CountTotal: true,
+	}
+	req := &types.QueryGetWithdrawalQueueRequest{
+		Pagination: pageReq,
+	}
+
+	// Query an empty queue
+	res, err := queryClient.GetWithdrawalQueue(context.Background(), req)
+	require.NoError(err)
+	require.Equal(0, len(res.Withdrawals), "expected no withdrawals in the queue yet")
+
+	// Prepare and add three withdrawals to the queue
+	delAddr = "story1hmjw3pvkjtndpg8wqppwdn8udd835qpan4hm0y"
+	valAddr = "storyvaloper1hmjw3pvkjtndpg8wqppwdn8udd835qpaa6r6y0"
+	evmAddr = common.HexToAddress("0x131D25EDE18178BAc9275b312001a63C081722d2")
+	withdrawals := []types.Withdrawal{
+		types.NewWithdrawal(1, delAddr, valAddr, evmAddr.String(), 100),
+		types.NewWithdrawal(2, delAddr, valAddr, evmAddr.String(), 200),
+		types.NewWithdrawal(3, delAddr, valAddr, evmAddr.String(), 300),
+	}
+	require.Len(withdrawals, 3)
+	for _, w := range withdrawals {
+		err = keeper.AddWithdrawalToQueue(ctx, w)
+		require.NoError(err)
+	}
+
+	// Query the first page of two withdrawals
+	res, err = queryClient.GetWithdrawalQueue(context.Background(), req)
+	require.NoError(err)
+	require.Equal(2, len(res.Withdrawals),
+		"expected 2 withdrawals after first page query, but found %d", len(res.Withdrawals))
+
+	// Query the next page for the remaining withdrawal
+	nextPage := res.Pagination.NextKey
+	require.NotNil(nextPage, "expected a next page key to be not nil")
+
+	pageReq.Key = nextPage
+	req = &types.QueryGetWithdrawalQueueRequest{
+		Pagination: pageReq,
+	}
+	res, err = queryClient.GetWithdrawalQueue(context.Background(), req)
+	require.NoError(err)
+	require.Equal(1, len(res.Withdrawals), "expected 1 withdrawal after second page query, but found %d", len(res.Withdrawals))
+}
diff --git a/client/x/evmstaking/keeper/keeper_test.go b/client/x/evmstaking/keeper/keeper_test.go
index 58c1018a..aee98e5b 100644
--- a/client/x/evmstaking/keeper/keeper_test.go
+++ b/client/x/evmstaking/keeper/keeper_test.go
@@ -15,6 +15,7 @@ import (
 	"github.com/cometbft/cometbft/crypto"
 	cmtproto "github.com/cometbft/cometbft/proto/tendermint/types"
 	dbm "github.com/cosmos/cosmos-db"
+	"github.com/cosmos/cosmos-sdk/baseapp"
 	"github.com/cosmos/cosmos-sdk/codec"
 	"github.com/cosmos/cosmos-sdk/codec/address"
 	codectypes "github.com/cosmos/cosmos-sdk/codec/types"
@@ -61,6 +62,7 @@ type TestSuite struct {
 	SlashingKeeper   *estestutil.MockSlashingKeeper
 	EVMStakingKeeper *keeper.Keeper
 	msgServer        types.MsgServiceServer
+	queryClient      types.QueryClient
 
 	encCfg moduletestutil.TestEncodingConfig
 }
@@ -142,6 +144,9 @@ func (s *TestSuite) SetupTest() {
 	s.Require().NoError(evmstakingKeeper.SetParams(s.Ctx, types.DefaultParams()))
 	s.EVMStakingKeeper = evmstakingKeeper
 	s.msgServer = keeper.NewMsgServerImpl(evmstakingKeeper)
+	queryHelper := baseapp.NewQueryServerTestHelper(s.Ctx, s.encCfg.InterfaceRegistry)
+	types.RegisterQueryServer(queryHelper, evmstakingKeeper)
+	s.queryClient = types.NewQueryClient(queryHelper)
 }
 
 func (s *TestSuite) TestLogger() {