simple unit tests for prepare proposal handler

x/lastlook/abci/preblock/preblock.go

@@ -93,7 +93,7 @@ func (h *LastLookPreBlockHandler) PreBlocker() sdk.PreBlocker {
 			return &sdk.ResponsePreBlock{}, nil
 		}
 
-		if err := h.keeper.StoreBatch(ctx, sdk.AccAddress(txBlob.Proposer), txBlob.Txs); err != nil {
+		if err := h.keeper.StoreBatch(ctx, ctx.BlockHeight()+1, sdk.AccAddress(txBlob.Proposer), txBlob.Txs); err != nil {
 			h.logger.Error("failed to store txs", "err", err)
 			return nil, err
 		}

x/lastlook/abci/proposals/preblock_test.go

@@ -0,0 +1,140 @@
+package proposals_test
+
+import (
+	"crypto/rand"
+	"testing"
+
+	"github.com/cometbft/cometbft/abci/types"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/stretchr/testify/require"
+
+	"github.com/neutron-org/neutron/v4/testutil/lastlook/keeper"
+	"github.com/neutron-org/neutron/v4/x/lastlook/abci/proposals"
+	lastlookkeeper "github.com/neutron-org/neutron/v4/x/lastlook/keeper"
+	lastlooktypes "github.com/neutron-org/neutron/v4/x/lastlook/types"
+)
+
+func randomTxsWithSize(n int, size int) [][]byte {
+	txs := make([][]byte, 0, n)
+	for i := 0; i < n; i++ {
+		buf := make([]byte, size)
+		if _, err := rand.Read(buf); err != nil {
+			panic(err)
+		}
+
+		txs = append(txs, buf)
+	}
+
+	return txs
+}
+
+func TestPrepareProposalHandler(t *testing.T) {
+	type testCase struct {
+		name string
+		// this method must return RequestPrepareProposal as the first argument, expected txs to be returned after proposal preparation and expected error
+		malleate func(*testing.T, sdk.Context, *lastlookkeeper.Keeper) (types.RequestPrepareProposal, [][]byte, error)
+	}
+
+	testCases := []testCase{
+		{
+			name: "All good, 10 txs in mempool, 0 in queue, height 0",
+			malleate: func(t *testing.T, ctx sdk.Context, k *lastlookkeeper.Keeper) (types.RequestPrepareProposal, [][]byte, error) {
+				txReq := randomTxsWithSize(10, 1)
+				expectedBatch := lastlooktypes.Batch{
+					Proposer: []byte("ProposerAddress"),
+					Txs:      txReq,
+				}
+				expectedBatchBz, err := expectedBatch.Marshal()
+				require.NoError(t, err)
+
+				return types.RequestPrepareProposal{
+						MaxTxBytes:      100,
+						Txs:             txReq,
+						Height:          0,
+						ProposerAddress: []byte("ProposerAddress"),
+					},
+					[][]byte{expectedBatchBz},
+					nil
+			},
+		},
+		{
+			name: "All good, 10 txs in mempool, 5 in queue, height 2",
+			malleate: func(t *testing.T, ctx sdk.Context, k *lastlookkeeper.Keeper) (types.RequestPrepareProposal, [][]byte, error) {
+				txsInQueue := randomTxsWithSize(5, 2)
+				require.NoError(t, k.StoreBatch(ctx, 2, sdk.AccAddress([]byte("ProposerAddress")), txsInQueue))
+
+				txReq := randomTxsWithSize(10, 1)
+				expectedBatch := lastlooktypes.Batch{
+					Proposer: []byte("ProposerAddress"),
+					Txs:      txReq,
+				}
+				expectedBatchBz, err := expectedBatch.Marshal()
+				require.NoError(t, err)
+
+				expectedTxs := make([][]byte, 0)
+				expectedTxs = append(expectedTxs, expectedBatchBz)
+				expectedTxs = append(expectedTxs, txsInQueue...)
+
+				return types.RequestPrepareProposal{
+						MaxTxBytes:      10000,
+						Txs:             txReq,
+						Height:          2,
+						ProposerAddress: []byte("ProposerAddress"),
+					},
+					expectedTxs,
+					nil
+			},
+		},
+		{
+			name: "All good, 100 txs in mempool each 1 byte (total 100 byte), 0 in queue, MaxTxSize limit 50, height 0",
+			malleate: func(t *testing.T, ctx sdk.Context, k *lastlookkeeper.Keeper) (types.RequestPrepareProposal, [][]byte, error) {
+				txReq := randomTxsWithSize(100, 1)
+				expectedBatch := lastlooktypes.Batch{
+					Proposer: []byte("ProposerAddress"),
+					// we expected only 11 txs to be in a batch because there is:
+					// * overhead for encoded Batch structure - 17 bytes;
+					// * overhead for each tx item in a batch - 3 bytes;
+					// * each tx consumes 1 byte of memory (just for simple calculations).
+					// In total with MaxTxBytes equals to 50 bytes, batch with 11 txs consumes exactly 50 bytes:
+					// 17 + 3*11 = 50
+					Txs: txReq[:11],
+				}
+				expectedBatchBz, err := expectedBatch.Marshal()
+				require.NoError(t, err)
+
+				expectedTxs := make([][]byte, 0)
+				expectedTxs = append(expectedTxs, expectedBatchBz)
+
+				return types.RequestPrepareProposal{
+						MaxTxBytes:      50,
+						Txs:             txReq,
+						Height:          0,
+						ProposerAddress: []byte("ProposerAddress"),
+					},
+					expectedTxs,
+					nil
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		lastLookKeeper, ctx := keeper.LastLookKeeper(t)
+		proposalHandler := proposals.NewProposalHandler(ctx.Logger(), lastLookKeeper)
+
+		t.Run(tc.name, func(t *testing.T) {
+			req, expectedTxs, expectedError := tc.malleate(t, ctx, lastLookKeeper)
+
+			resp, err := proposalHandler.PrepareProposalHandler()(ctx, &req)
+			require.Equal(t, expectedError, err)
+			require.Equal(t, expectedTxs, resp.Txs)
+
+			totalBytes := int64(0)
+			for _, tx := range resp.Txs {
+				totalBytes += int64(len(tx))
+			}
+
+			// total bytes consumption must not exceed provided req.MaxTxBytes
+			require.GreaterOrEqual(t, req.MaxTxBytes, totalBytes)
+		})
+	}
+}

x/lastlook/abci/proposals/proposals.go

@@ -57,7 +57,7 @@ func (h *ProposalHandler) PrepareProposalHandler() sdk.PrepareProposalHandler {
 
 		consumedBytes := int64(0)
 		// increase consumedBytes counter by a size of all transaction from a batch that must be inserted to a block
-		for _, tx := range req.Txs {
+		for _, tx := range currentBatch.Txs {
 			consumedBytes += int64(len(tx))
 		}
 
@@ -81,14 +81,15 @@ func (h *ProposalHandler) PrepareProposalHandler() sdk.PrepareProposalHandler {
 
 			newBatch.Txs = append(newBatch.Txs, mempoolTx)
 
+			batchSize := int64(newBatch.XXX_Size())
 			// if final size of a serialised tx is too big to insert into a block, remove last inserted tx and continue
 			// iterating over txs from mempool, maybe there is some small tx we can still fit in
-			if int64(newBatch.XXX_Size())+consumedBytes > req.MaxTxBytes {
+			if batchSize+consumedBytes > req.MaxTxBytes {
 				newBatch.Txs = newBatch.Txs[:len(newBatch.Txs)-1]
 			}
 		}
 
-		newBatchBz, err := h.lastlookKeeper.GetCodec().Marshal(&newBatch)
+		newBatchBz, err := newBatch.Marshal()
 		if err != nil {
 			h.logger.Error("failed to marshal txs batch", "err", err)
 			return nil, err

x/lastlook/keeper/queue.go

@@ -8,7 +8,7 @@ import (
 )
 
 // StoreBatch stores a slice of txs and an address of a proposer who proposed a block into a queue
-func (k *Keeper) StoreBatch(ctx sdk.Context, proposer sdk.Address, txs [][]byte) error {
+func (k *Keeper) StoreBatch(ctx sdk.Context, height int64, proposer sdk.Address, txs [][]byte) error {
 	store := ctx.KVStore(k.storeKey)
 	blob := types.Batch{
 		Proposer: proposer.Bytes(),
@@ -20,7 +20,7 @@ func (k *Keeper) StoreBatch(ctx sdk.Context, proposer sdk.Address, txs [][]byte)
 		return errors.Wrapf(types.ErrProtoMarshal, "failed to marshal txs blob: %v", err)
 	}
 
-	store.Set(types.GetTxsQueueKey(ctx.BlockHeight()+1), bz)
+	store.Set(types.GetTxsQueueKey(height), bz)
 
 	return nil
 }
@@ -31,17 +31,17 @@ func (k *Keeper) GetBatch(ctx sdk.Context, blockHeight int64) (*types.Batch, err
 
 	var blob types.Batch
 
-	if blockHeight == 1 {
+	if blockHeight <= 1 {
 		return &blob, nil
 	}
 
 	bz := store.Get(types.GetTxsQueueKey(blockHeight))
 	if bz == nil {
-		return nil, errors.Wrapf(types.ErrNoBlob, "no txs blob found for a block %d", blockHeight)
+		return nil, errors.Wrapf(types.ErrNoBlob, "no txs batch found for a block %d", blockHeight)
 	}
 
 	if err := k.cdc.Unmarshal(bz, &blob); err != nil {
-		return nil, errors.Wrapf(types.ErrProtoUnmarshal, "failed to unmarshal tx blob: %v", err)
+		return nil, errors.Wrapf(types.ErrProtoUnmarshal, "failed to unmarshal tx batch: %v", err)
 	}
 
 	return &blob, nil