test: Fix mempool tests

b-harvest · Jan 10, 2025 · 88ce42e · 88ce42e
1 parent a669614
commit 88ce42e
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 48 deletions.
diff --git a/mempool/clist_mempool.go b/mempool/clist_mempool.go
@@ -257,29 +257,27 @@ func (mem *CListMempool) checkTxAsync(tx types.Tx, txInfo TxInfo, prepareCb func
 			mem.updateMtx.RUnlock()
 			panic(r)
 		}
+		mem.updateMtx.RUnlock()
 	}()
 
 	err := mem.prepareCheckTx(tx, txInfo)
 	if prepareCb != nil {
 		prepareCb(err)
 	}
 	if err != nil {
-		mem.updateMtx.RUnlock()
 		return
 	}
 
 	// CONTRACT: `app.CheckTxAsync()` should check whether `GasWanted` is valid (0 <= GasWanted <= block.masGas)
 	reqRes, err := mem.proxyAppConn.CheckTxAsync(context.TODO(), &abci.RequestCheckTx{Tx: tx})
 	if err != nil {
-		mem.updateMtx.RUnlock()
 		return
 	}
 	reqRes.SetCallback(func(res *abci.Response) {
 		mem.reqResCb(tx, txInfo, res, func(response *abci.Response) {
 			if checkTxCb != nil {
 				checkTxCb(response)
 			}
-			mem.updateMtx.RUnlock()
 		})
 	})
 }

diff --git a/mempool/clist_mempool_test.go b/mempool/clist_mempool_test.go
@@ -292,17 +292,12 @@ func TestMempoolUpdateDoesNotPanicWhenApplicationMissedTx(t *testing.T) {
 	txs := []types.Tx{[]byte{0x01}, []byte{0x02}, []byte{0x03}, []byte{0x04}}
 	for _, tx := range txs {
 		reqRes := newReqRes(tx, abci.CodeTypeOK, abci.CheckTxType_New)
-		mockClient.On("CheckTxAsync", mock.Anything, mock.Anything, mock.Anything).Run(
-			func(args mock.Arguments) {
-				callback := args[2].(abciclient.CheckTxCallback)
-				callback(&abci.Response{
-					Value: &abci.Response_CheckTx{
-						CheckTx: &abci.ResponseCheckTx{
-							Code: abci.CodeTypeOK,
-						},
-					},
-				})
-			}).Return(reqRes, nil)
+		mockClient.On("CheckTxAsync", mock.Anything, mock.Anything).Return(
+			func(ctx context.Context, req *abci.RequestCheckTx) (*abciclient.ReqRes, error) {
+				reqRes.SetDone()
+				return reqRes, nil
+			},
+		)
 		mp.CheckTxAsync(tx, TxInfo{}, nil, nil)
 
 		// ensure that the callback that the mempool sets on the ReqRes is run.
@@ -848,17 +843,12 @@ func TestMempoolSyncRecheckTxReturnError(t *testing.T) {
 	txs := []types.Tx{[]byte{0x01}, []byte{0x02}}
 	for _, tx := range txs {
 		reqRes := newReqRes(tx, abci.CodeTypeOK, abci.CheckTxType_New)
-		mockClient.On("CheckTxAsync", mock.Anything, mock.Anything, mock.Anything).Run(
-			func(args mock.Arguments) {
-				callback := args[2].(abciclient.CheckTxCallback)
-				callback(&abci.Response{
-					Value: &abci.Response_CheckTx{
-						CheckTx: &abci.ResponseCheckTx{
-							Code: abci.CodeTypeOK,
-						},
-					},
-				})
-			}).Return(reqRes, nil).Once()
+		mockClient.On("CheckTxAsync", mock.Anything, mock.Anything).Return(
+			func(ctx context.Context, req *abci.RequestCheckTx) (*abciclient.ReqRes, error) {
+				reqRes.SetDone()
+				return reqRes, nil
+			},
+		).Once()
 		mp.CheckTxAsync(tx, TxInfo{}, nil, nil)
 
 		// ensure that the callback that the mempool sets on the ReqRes is run.
@@ -902,41 +892,32 @@ func TestMempoolAsyncRecheckTxReturnError(t *testing.T) {
 	txs := []types.Tx{[]byte{0x01}, []byte{0x02}, []byte{0x03}, []byte{0x04}}
 	for _, tx := range txs {
 		reqRes := newReqRes(tx, abci.CodeTypeOK, abci.CheckTxType_New)
-		mockClient.On("CheckTxAsync", mock.Anything, mock.Anything, mock.Anything).Run(
-			func(args mock.Arguments) {
-				callback := args[2].(abciclient.CheckTxCallback)
-				callback(&abci.Response{
-					Value: &abci.Response_CheckTx{
-						CheckTx: &abci.ResponseCheckTx{
-							Code: abci.CodeTypeOK,
-						},
-					},
-				})
-			}).Return(reqRes, nil).Once()
+		mockClient.On("CheckTxAsync", mock.Anything, mock.Anything).Return(
+			func(ctx context.Context, req *abci.RequestCheckTx) (*abciclient.ReqRes, error) {
+				reqRes.SetDone()
+				return reqRes, nil
+			},
+		).Once()
 		mp.CheckTxAsync(tx, TxInfo{}, nil, nil)
 
 		// ensure that the callback that the mempool sets on the ReqRes is run.
 		reqRes.InvokeCallback()
 	}
-	time.Sleep(20 * time.Millisecond)
+	time.Sleep(100 * time.Millisecond)
 	// The 4 txs are added to the mempool.
 	require.Len(t, txs, mp.Size())
 
 	// Check that recheck has not started.
 	mockClient.AssertExpectations(t)
 
 	// One call to CheckTxAsync per tx, for rechecking.
-	mockClient.On("CheckTxAsync", mock.Anything, mock.Anything, mock.Anything).Run(
-		func(args mock.Arguments) {
-			callback := args[2].(abciclient.CheckTxCallback)
-			callback(&abci.Response{
-				Value: &abci.Response_CheckTx{
-					CheckTx: &abci.ResponseCheckTx{
-						Code: abci.CodeTypeOK,
-					},
-				},
-			})
-		}).Return(nil, nil).Times(4)
+	mockClient.On("CheckTxAsync", mock.Anything, mock.Anything).Return(
+		func(ctx context.Context, req *abci.RequestCheckTx) (*abciclient.ReqRes, error) {
+			reqRes := newReqRes(req.Tx, abci.CodeTypeOK, abci.CheckTxType_Recheck)
+			reqRes.SetDone()
+			return reqRes, nil
+		},
+	).Times(4)
 
 	// On the async client, the callbacks are executed when flushing the connection. The app replies
 	// to the request for the first tx (valid) and for the third tx (invalid), so the callback is