From d82fb33d1ecede4bddb50b54a444131e0bf26cda Mon Sep 17 00:00:00 2001
From: Eric Warehime <eric.warehime@gmail.com>
Date: Tue, 16 Jan 2024 13:14:20 -0800
Subject: [PATCH] test: mempool select test (#360)

---
 block/base/mempool_test.go   | 103 +++++++++++++++++++++++++++++++++++
 testutils/mempool/mempool.go |   1 +
 2 files changed, 104 insertions(+)

diff --git a/block/base/mempool_test.go b/block/base/mempool_test.go
index c02f610..36eb076 100644
--- a/block/base/mempool_test.go
+++ b/block/base/mempool_test.go
@@ -13,6 +13,12 @@ import (
 	"github.com/skip-mev/block-sdk/testutils"
 )
 
+type txGen struct {
+	acc    testutils.Account
+	nonce  uint64
+	amount sdk.Coin
+}
+
 func TestMempoolComparison(t *testing.T) {
 	acct := testutils.RandomAccounts(rand.New(rand.NewSource(1)), 2)
 	txc := testutils.CreateTestEncodingConfig().TxConfig
@@ -87,3 +93,100 @@ func TestMempoolComparison(t *testing.T) {
 		require.Equal(t, 0, output)
 	})
 }
+
+func TestMempoolSelect(t *testing.T) {
+	acct := testutils.RandomAccounts(rand.New(rand.NewSource(1)), 2)
+	txc := testutils.CreateTestEncodingConfig().TxConfig
+	ctx := testutils.CreateBaseSDKContext(t)
+	se := signerextraction.NewDefaultAdapter()
+	mp := NewMempool(
+		DefaultTxPriority(),
+		txc.TxEncoder(),
+		se,
+		1000,
+	)
+	tests := []struct {
+		name     string
+		inputs   []txGen
+		expected []txGen
+	}{
+		{
+			name: "test1",
+			inputs: []txGen{
+				{
+					acc:    acct[0],
+					amount: sdk.NewCoin("stake", sdkmath.NewInt(1)),
+					nonce:  0,
+				},
+				{
+					acc:    acct[1],
+					amount: sdk.NewCoin("notstake", sdkmath.NewInt(1)),
+					nonce:  0,
+				},
+				{
+					acc:    acct[0],
+					amount: sdk.NewCoin("notstake", sdkmath.NewInt(2)),
+					nonce:  1,
+				},
+				{
+					acc:    acct[1],
+					amount: sdk.NewCoin("stake", sdkmath.NewInt(2)),
+					nonce:  1,
+				},
+			},
+			expected: []txGen{
+				{
+					acc:    acct[1],
+					amount: sdk.NewCoin("notstake", sdkmath.NewInt(1)),
+					nonce:  0,
+				},
+				{
+					acc:    acct[1],
+					amount: sdk.NewCoin("stake", sdkmath.NewInt(2)),
+					nonce:  1,
+				},
+				{
+					acc:    acct[0],
+					amount: sdk.NewCoin("stake", sdkmath.NewInt(1)),
+					nonce:  0,
+				},
+				{
+					acc:    acct[0],
+					amount: sdk.NewCoin("notstake", sdkmath.NewInt(2)),
+					nonce:  1,
+				},
+			},
+		},
+	}
+	for _, tc := range tests {
+		// insert all input txs
+		t.Run(tc.name, func(t *testing.T) {
+			for _, tx := range tc.inputs {
+				inputTx, err := testutils.CreateTx(txc, tx.acc, tx.nonce, 0, nil, tx.amount)
+				require.NoError(t, err)
+				err = mp.Insert(ctx, inputTx)
+				require.NoError(t, err)
+			}
+		})
+		// extract all txs via select
+		var output []sdk.Tx
+		for iter := mp.Select(ctx, nil); iter != nil; iter = iter.Next() {
+			output = append(output, iter.Tx())
+		}
+		// validate the order matches the expected order
+		require.Equal(t, len(tc.expected), len(output))
+		for i, tx := range output {
+
+			sigs, err := se.GetSigners(tx)
+			require.NoError(t, err)
+			require.Equal(t, 1, len(sigs))
+			require.Equal(t, tc.expected[i].acc.Address, sigs[0].Signer)
+			feeTx, ok := tx.(sdk.FeeTx)
+			require.True(t, ok)
+			require.Equal(t, tc.expected[i].amount.Denom, feeTx.GetFee().GetDenomByIndex(0))
+			require.Equal(t, tc.expected[i].amount.Amount, feeTx.GetFee().AmountOf(feeTx.GetFee().GetDenomByIndex(0)))
+			require.Equal(t, tc.expected[i].nonce, sigs[0].Sequence)
+
+		}
+	}
+}
diff --git a/testutils/mempool/mempool.go b/testutils/mempool/mempool.go
index edfffd1..b6bb41e 100644
--- a/testutils/mempool/mempool.go
+++ b/testutils/mempool/mempool.go
@@ -3,6 +3,7 @@ package mempool
 import (
 	"cosmossdk.io/log"
 	"cosmossdk.io/math"
+
 	"github.com/skip-mev/block-sdk/block/mocks"
 
 	signerextraction "github.com/skip-mev/block-sdk/adapters/signer_extraction_adapter"