test: add test file for fetch.go (#1401)

* test: add test file for fetch.go

Signed-off-by: Ji Hwan <jkim@polygon.technology>

* chore: lint

Signed-off-by: Ji Hwan <jkim@polygon.technology>

* chore: cleanup

Signed-off-by: Ji Hwan <jkim@polygon.technology>

* docs: add docs in test code

Signed-off-by: Ji Hwan <jkim@polygon.technology>

* test: generate mock for pool interface

Signed-off-by: Ji Hwan <jkim@polygon.technology>

* test: add test for TxPool's OnNewBlock function

Signed-off-by: Ji Hwan <jkim@polygon.technology>

---------

Signed-off-by: Ji Hwan <jkim@polygon.technology>
Co-authored-by: Valentin Staykov <79150443+V-Staykov@users.noreply.github.com>

zk/txpool/fetch_test.go

@@ -0,0 +1,96 @@
+package txpool
+
+import (
+	"context"
+	"encoding/hex"
+	"fmt"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/ledgerwatch/erigon-lib/common/datadir"
+	"github.com/ledgerwatch/erigon-lib/common/u256"
+	"github.com/ledgerwatch/erigon-lib/direct"
+	"github.com/ledgerwatch/erigon-lib/gointerfaces"
+	"github.com/ledgerwatch/erigon-lib/gointerfaces/remote"
+	"github.com/ledgerwatch/erigon-lib/gointerfaces/sentry"
+	"github.com/ledgerwatch/erigon-lib/kv/kvcache"
+	"github.com/ledgerwatch/erigon-lib/kv/memdb"
+	"github.com/ledgerwatch/erigon-lib/kv/temporal/temporaltest"
+	"github.com/ledgerwatch/erigon-lib/txpool"
+	"github.com/ledgerwatch/erigon-lib/txpool/txpoolcfg"
+	types "github.com/ledgerwatch/erigon-lib/types"
+	"github.com/ledgerwatch/erigon/eth/ethconfig"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/mock/gomock"
+)
+
+var peerID types.PeerID = gointerfaces.ConvertHashToH512([64]byte{0x12, 0x34, 0x50}) // "12345"
+
+// DecodeHex converts a hex string to a byte array.
+func decodeHex(in string) []byte {
+	payload, err := hex.DecodeString(in)
+	if err != nil {
+		panic(err)
+	}
+	return payload
+}
+
+func TestFetch(t *testing.T) {
+	assert, require := assert.New(t), require.New(t)
+	ch := make(chan types.Announcements, 100)
+	_, coreDB, _ := temporaltest.NewTestDB(t, datadir.New(t.TempDir()))
+	defer coreDB.Close()
+	db := memdb.NewTestPoolDB(t)
+	path := fmt.Sprintf("/tmp/db-test-%v", time.Now().UTC().Format(time.RFC3339Nano))
+	txPoolDB := newTestTxPoolDB(t, path)
+	defer txPoolDB.Close()
+	aclsDB := newTestACLDB(t, path)
+	defer aclsDB.Close()
+
+	// Check if the dbs are created.
+	require.NotNil(t, db)
+	require.NotNil(t, txPoolDB)
+	require.NotNil(t, aclsDB)
+
+	cfg := txpoolcfg.DefaultConfig
+	ethCfg := &ethconfig.Defaults
+	sendersCache := kvcache.New(kvcache.DefaultCoherentConfig)
+	pool, err := New(ch, coreDB, cfg, ethCfg, sendersCache, *u256.N1, nil, nil, aclsDB)
+	assert.NoError(err)
+	require.True(pool != nil)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	ctrl := gomock.NewController(t)
+	remoteKvClient := remote.NewMockKVClient(ctrl)
+	sentryServer := sentry.NewMockSentryServer(ctrl)
+
+	m := txpool.NewMockSentry(ctx, sentryServer)
+	sentryClient := direct.NewSentryClientDirect(direct.ETH66, m)
+	fetch := NewFetch(ctx, []direct.SentryClient{sentryClient}, pool, remoteKvClient, nil, nil, *u256.N1)
+	var wg sync.WaitGroup
+	fetch.SetWaitGroup(&wg)
+	// The corresponding WaitGroup.Done() will be called by the Sentry.
+	// First will be called by (txpool.MockSentry).Messages
+	// Second will be called by (txpool.MockSentry).PeerEvents
+	m.StreamWg.Add(2)
+	fetch.ConnectSentries()
+	m.StreamWg.Wait()
+
+	// Send one transaction id with ETH66 protocol.
+	// The corresponding WaitGroup.Done() will be called by the fetch.receiveMessage()
+	wg.Add(1)
+	errs := m.Send(&sentry.InboundMessage{
+		Id:     sentry.MessageId_NEW_POOLED_TRANSACTION_HASHES_66,
+		Data:   decodeHex("e1a0595e27a835cd79729ff1eeacec3120eeb6ed1464a04ec727aaca734ead961328"),
+		PeerId: peerID,
+	})
+	for i, err := range errs {
+		if err != nil {
+			t.Errorf("sending new pool txn hashes 66 (%d): %v", i, err)
+		}
+	}
+	wg.Wait()
+}

zk/txpool/pool_test.go

@@ -3,6 +3,7 @@ package txpool
 import (
 	"context"
 	"fmt"
+	"io"
 	"testing"
 	"time"
 
@@ -12,6 +13,7 @@ import (
 	"github.com/ledgerwatch/erigon-lib/common/u256"
 	"github.com/ledgerwatch/erigon-lib/gointerfaces"
 	"github.com/ledgerwatch/erigon-lib/gointerfaces/remote"
+	"github.com/ledgerwatch/erigon-lib/kv"
 	"github.com/ledgerwatch/erigon-lib/kv/kvcache"
 	"github.com/ledgerwatch/erigon-lib/kv/memdb"
 	"github.com/ledgerwatch/erigon-lib/kv/temporal/temporaltest"
@@ -20,6 +22,8 @@ import (
 	"github.com/ledgerwatch/erigon/eth/ethconfig"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.uber.org/mock/gomock"
+	"google.golang.org/grpc"
 )
 
 func TestNonceFromAddress(t *testing.T) {
@@ -169,3 +173,74 @@ func TestNonceFromAddress(t *testing.T) {
 		}
 	}
 }
+
+func TestOnNewBlock(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	coreDB, db := memdb.NewTestDB(t), memdb.NewTestDB(t)
+	ctrl := gomock.NewController(t)
+
+	stream := remote.NewMockKV_StateChangesClient(ctrl)
+	i := 0
+	stream.EXPECT().
+		Recv().
+		DoAndReturn(func() (*remote.StateChangeBatch, error) {
+			if i > 0 {
+				return nil, io.EOF
+			}
+			i++
+			return &remote.StateChangeBatch{
+				StateVersionId: 1,
+				ChangeBatch: []*remote.StateChange{
+					{
+						Txs: [][]byte{
+							decodeHex(types.TxParseMainnetTests[0].PayloadStr),
+							decodeHex(types.TxParseMainnetTests[1].PayloadStr),
+							decodeHex(types.TxParseMainnetTests[2].PayloadStr),
+						},
+						BlockHeight: 1,
+						BlockHash:   gointerfaces.ConvertHashToH256([32]byte{}),
+					},
+				},
+			}, nil
+		}).
+		AnyTimes()
+
+	stateChanges := remote.NewMockKVClient(ctrl)
+	stateChanges.
+		EXPECT().
+		StateChanges(gomock.Any(), gomock.Any(), gomock.Any()).
+		DoAndReturn(func(_ context.Context, _ *remote.StateChangeRequest, _ ...grpc.CallOption) (remote.KV_StateChangesClient, error) {
+			return stream, nil
+		})
+
+	pool := NewMockPool(ctrl)
+	pool.EXPECT().
+		ValidateSerializedTxn(gomock.Any()).
+		DoAndReturn(func(_ []byte) error {
+			return nil
+		}).
+		Times(3)
+
+	var minedTxs types.TxSlots
+	pool.EXPECT().
+		OnNewBlock(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
+		DoAndReturn(
+			func(
+				_ context.Context,
+				_ *remote.StateChangeBatch,
+				_ types.TxSlots,
+				minedTxsArg types.TxSlots,
+				_ kv.Tx,
+			) error {
+				minedTxs = minedTxsArg
+				return nil
+			},
+		).
+		Times(1)
+
+	fetch := NewFetch(ctx, nil, pool, stateChanges, coreDB, db, *u256.N1)
+	err := fetch.handleStateChanges(ctx, stateChanges)
+	assert.ErrorIs(t, io.EOF, err)
+	assert.Equal(t, 3, len(minedTxs.Txs))
+}