chain: patch unit tests for batchGetRawTxes

chain/mempool_test.go

@@ -4,6 +4,7 @@ import (
 	"errors"
 	"math"
 	"testing"
+	"time"
 
 	"github.com/btcsuite/btcd/btcutil"
 	"github.com/btcsuite/btcd/chaincfg/chainhash"
@@ -12,6 +13,10 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+// testTimeout is used to ensure the test will time out if the expected result
+// is not returned in 5 seconds.
+const testTimeout = 5 * time.Second
+
 // TestCachedInputs tests that the cachedInputs works as expected.
 func TestCachedInputs(t *testing.T) {
 	require := require.New(t)
@@ -655,3 +660,236 @@ func TestGetRawTxIgnoreErr(t *testing.T) {
 	// Assert the mock client was called as expected.
 	mockReceiver.AssertExpectations(t)
 }
+
+// TestBatchGetRawTxesOnBatchSize checks that the batch size is properly
+// handled. It defines a testing batch size of 3, and creates 7 testing
+// transactions. Then it asserts there are 3 batches created and handled.
+func TestBatchGetRawTxesOnBatchSize(t *testing.T) {
+	require := require.New(t)
+
+	const (
+		// Define a small batch size for testing only.
+		testBatchSize = 3
+
+		// Create 7 test transactions so we can hit our batching logic
+		// - we should create two full batches and one batch with the
+		// remaining transaction.
+		numTxes = testBatchSize*2 + 1
+	)
+
+	// Create a mock client and init our mempool.
+	mockRPC := &mockRPCClient{}
+	m := newMempool(&mempoolConfig{
+		client:            mockRPC,
+		batchWaitInterval: 0,
+		getRawTxBatchSize: testBatchSize,
+	})
+
+	// Create test transactions and mempool state.
+	mempool := make([]*chainhash.Hash, 0, numTxes)
+
+	// Create a map of raw tx response receivers, keyed by txid.
+	mockTxResponses := make(map[chainhash.Hash]*btcutil.Tx, numTxes)
+
+	// Fill up the slices and mock the methods.
+	for i := 0; i < numTxes; i++ {
+		// Create testing transactions.
+		op := wire.OutPoint{Hash: chainhash.Hash{byte(i)}}
+		tx := &wire.MsgTx{
+			LockTime: 1,
+			TxIn:     []*wire.TxIn{{PreviousOutPoint: op}},
+		}
+
+		// Fill the testing mempool.
+		txHash := tx.TxHash()
+		mempool = append(mempool, &txHash)
+
+		// Create a testing resposne receiver to be returned by
+		// GetRawTransactionAsync.
+		mockTxReceiver := make(rpcclient.FutureGetRawTransactionResult)
+
+		// Add this tx to our mocked responses.
+		btcTx := btcutil.NewTx(tx)
+		mockTxResponses[txHash] = btcTx
+
+		// Mock `GetRawTransactionAsync` to return the mocked value.
+		mockRPC.On("GetRawTransactionAsync",
+			&txHash).Return(mockTxReceiver).Once()
+	}
+
+	// Mock the rawTxReceiver to find and return the tx found in map
+	// `mockTxResponses`.
+	m.cfg.rawTxReceiver = func(txid chainhash.Hash,
+		reciever getRawTxReceiver) *btcutil.Tx {
+
+		btcTx, ok := mockTxResponses[txid]
+		require.Truef(ok, "unexpected receiver for %v", txid)
+
+		return btcTx
+	}
+
+	// We expect to send the batched requests three times - two for the
+	// full batch and one for the remaining batch.
+	mockRPC.On("Send").Return(nil).Times(3)
+
+	// Call the method under test.
+	newTxes, err := m.batchGetRawTxes(mempool, true)
+	require.NoError(err)
+
+	// Validate we have the expected number of transactions returned.
+	require.Len(newTxes, numTxes)
+
+	// Assert the mock methods are called as expected.
+	mockRPC.AssertExpectations(t)
+}
+
+// TestBatchGetRawTxesOnShutdown checks that the method returns immediately
+// when the mempool is shutting down.
+func TestBatchGetRawTxesOnShutdown(t *testing.T) {
+	require := require.New(t)
+
+	// Create a mock client and init our mempool.
+	mockRPC := &mockRPCClient{}
+	m := newMempool(&mempoolConfig{
+		client:            mockRPC,
+		batchWaitInterval: 0,
+		getRawTxBatchSize: 1,
+	})
+
+	// Create a normal transaction.
+	op1 := wire.OutPoint{Hash: chainhash.Hash{1}}
+	tx1 := &wire.MsgTx{
+		LockTime: 1,
+		TxIn: []*wire.TxIn{
+			{PreviousOutPoint: op1},
+		},
+	}
+	tx1Hash := tx1.TxHash()
+
+	// Create the current mempool state.
+	mempool := []*chainhash.Hash{&tx1Hash}
+
+	// Shutdown the mempool before call the method.
+	m.Shutdown()
+
+	// Call the method under test.
+	newTxes, err := m.batchGetRawTxes(mempool, true)
+
+	// We expect no error and transactions.
+	require.NoError(err)
+	require.Empty(newTxes)
+
+	// Assert GetRawTransaction is not called because mempool has quit.
+	mockRPC.AssertNotCalled(t, "GetRawTransactionAsync")
+	mockRPC.AssertNotCalled(t, "Send")
+}
+
+// TestBatchGetRawTxesOnWait checks that the method stays on hold once the
+// first batch is finished.
+func TestBatchGetRawTxesOnWait(t *testing.T) {
+	require := require.New(t)
+
+	const (
+		// Define a long wait interval for testing only.
+		testWaitInterval = 10 * time.Minute
+
+		// Define a small batch size for testing only.
+		testBatchSize = 3
+
+		// Create 4 test transactions so we can hit our batching logic
+		// once and then starts waiting.
+		numTxes = testBatchSize + 1
+	)
+
+	// Create a mock client and init our mempool.
+	mockRPC := &mockRPCClient{}
+	m := newMempool(&mempoolConfig{
+		client:            mockRPC,
+		batchWaitInterval: testWaitInterval,
+		getRawTxBatchSize: testBatchSize,
+	})
+
+	// Create test transactions and mempool state.
+	mempool := make([]*chainhash.Hash, 0, numTxes)
+
+	// Create a map of raw tx response receivers, keyed by txid.
+	mockTxResponses := make(map[chainhash.Hash]*btcutil.Tx, numTxes)
+
+	// Fill up the slices.
+	for i := 0; i < numTxes; i++ {
+		// Create testing transactions.
+		op := wire.OutPoint{Hash: chainhash.Hash{byte(i)}}
+		tx := &wire.MsgTx{
+			LockTime: 1,
+			TxIn:     []*wire.TxIn{{PreviousOutPoint: op}},
+		}
+
+		// Fill the testing mempool.
+		txHash := tx.TxHash()
+		mempool = append(mempool, &txHash)
+
+		// Add this tx to our mocked responses.
+		btcTx := btcutil.NewTx(tx)
+		mockTxResponses[txHash] = btcTx
+	}
+
+	// Mock GetRawTransactionAsync. We expect it to be called 3 times.
+	for i := 0; i < testBatchSize; i++ {
+		// Create a testing resposne receiver.
+		mockTxReceiver := make(rpcclient.FutureGetRawTransactionResult)
+
+		// Mock `GetRawTransactionAsync` to return the mocked value.
+		mockRPC.On("GetRawTransactionAsync",
+			mempool[i]).Return(mockTxReceiver).Once()
+	}
+
+	// Mock the rawTxReceiver to find and return the tx found in map
+	// `mockTxResponses`.
+	m.cfg.rawTxReceiver = func(txid chainhash.Hash,
+		reciever getRawTxReceiver) *btcutil.Tx {
+
+		btcTx, ok := mockTxResponses[txid]
+		require.Truef(ok, "unexpected receiver for %v", txid)
+
+		return btcTx
+	}
+
+	// We expect to send the batched requests exactly one time as the
+	// second batch will be blocked on the waiting.
+	mockRPC.On("Send").Return(nil).Once()
+
+	var (
+		err     error
+		newTxes []*wire.MsgTx
+		done    = make(chan struct{})
+	)
+
+	// Call the method under test in a goroutine so we don't need to wait.
+	go func() {
+		newTxes, err = m.batchGetRawTxes(mempool, true)
+
+		// Signal it's returned.
+		close(done)
+	}()
+
+	// Sleep one second to allow the mempool moves to the point where the
+	// first batch is finished and it's now blocked on the waiting. We then
+	// shut down the mempool so batchGetRawTxes will return immediately.
+	time.Sleep(1 * time.Second)
+	m.Shutdown()
+
+	// Catch the returned values with timeout.
+	select {
+	case <-done:
+		// Assert no error is returned, and we should get a nil slice
+		// since mempool is shut down.
+		require.NoError(err)
+		require.Nil(newTxes)
+
+	case <-time.After(testTimeout):
+		require.Fail("timeout waiting for batchGetRawTxes")
+	}
+
+	// Assert the mock methods are called as expected.
+	mockRPC.AssertExpectations(t)
+}