From b8b25700f698b7d59d70e62901484b346c2ff331 Mon Sep 17 00:00:00 2001
From: Rootul Patel <rootulp@gmail.com>
Date: Wed, 25 Sep 2024 11:21:29 -0400
Subject: [PATCH] refactor: TestTxPool_ConcurrentTxs

---
 mempool/cat/pool_test.go | 60 ++++++++++++++++++++--------------------
 1 file changed, 30 insertions(+), 30 deletions(-)

diff --git a/mempool/cat/pool_test.go b/mempool/cat/pool_test.go
index 03e58d591..03d917dcd 100644
--- a/mempool/cat/pool_test.go
+++ b/mempool/cat/pool_test.go
@@ -480,19 +480,21 @@ func TestTxPool_CheckTxSamePeer(t *testing.T) {
 	require.Error(t, txmp.CheckTx(tx, nil, mempool.TxInfo{SenderID: peerID}))
 }
 
+// TestTxPool_ConcurrentTxs adds a bunch of txs to the txPool (via checkTx) and
+// then reaps transactions from the mempool. At the end it asserts that the
+// mempool is empty.
 func TestTxPool_ConcurrentTxs(t *testing.T) {
-	txmp := setup(t, 100)
-	rng := rand.New(rand.NewSource(time.Now().UnixNano()))
+	cacheSize := 10
+	txPool := setup(t, cacheSize)
 	checkTxDone := make(chan struct{})
 
 	var wg sync.WaitGroup
-
 	wg.Add(1)
 	go func() {
-		for i := 0; i < 20; i++ {
-			_ = checkTxs(t, txmp, 100, 0)
-			dur := rng.Intn(1000-500) + 500
-			time.Sleep(time.Duration(dur) * time.Millisecond)
+		for i := 0; i < 10; i++ {
+			numTxs := 10
+			peerID := uint16(0)
+			_ = checkTxs(t, txPool, numTxs, peerID)
 		}
 
 		wg.Done()
@@ -505,33 +507,18 @@ func TestTxPool_ConcurrentTxs(t *testing.T) {
 		defer ticker.Stop()
 		defer wg.Done()
 
-		var height int64 = 1
-
+		height := int64(1)
 		for range ticker.C {
-			reapedTxs := txmp.ReapMaxTxs(200)
+			reapedTxs := txPool.ReapMaxTxs(50)
 			if len(reapedTxs) > 0 {
-				responses := make([]*abci.ResponseDeliverTx, len(reapedTxs))
-				for i := 0; i < len(responses); i++ {
-					var code uint32
-
-					if i%10 == 0 {
-						code = 100
-					} else {
-						code = abci.CodeTypeOK
-					}
-
-					responses[i] = &abci.ResponseDeliverTx{Code: code}
-				}
-
-				txmp.Lock()
-				require.NoError(t, txmp.Update(height, reapedTxs, responses, nil, nil))
-				txmp.Unlock()
-
+				responses := generateResponses(len(reapedTxs))
+				err := txPool.Update(height, reapedTxs, responses, nil, nil)
+				require.NoError(t, err)
 				height++
 			} else {
-				// only return once we know we finished the CheckTx loop
 				select {
 				case <-checkTxDone:
+					// only return once we know we finished the CheckTx loop
 					return
 				default:
 				}
@@ -540,8 +527,21 @@ func TestTxPool_ConcurrentTxs(t *testing.T) {
 	}()
 
 	wg.Wait()
-	require.Zero(t, txmp.Size())
-	require.Zero(t, txmp.SizeBytes())
+	assert.Zero(t, txPool.Size())
+	assert.Zero(t, txPool.SizeBytes())
+}
+
+func generateResponses(numResponses int) (responses []*abci.ResponseDeliverTx) {
+	for i := 0; i < numResponses; i++ {
+		var response *abci.ResponseDeliverTx
+		if i%2 == 0 {
+			response = &abci.ResponseDeliverTx{Code: abci.CodeTypeOK}
+		} else {
+			response = &abci.ResponseDeliverTx{Code: 100}
+		}
+		responses = append(responses, response)
+	}
+	return responses
 }
 
 func TestTxPool_ExpiredTxs_Timestamp(t *testing.T) {