polygon/heimdall: unit test entityFetcher.FetchAllEntities (#10379)

polygon/heimdall/entity_fetcher.go

@@ -13,14 +13,17 @@ import (
 type entityFetcher[TEntity Entity] interface {
 	FetchLastEntityId(ctx context.Context) (uint64, error)
 	FetchEntitiesRange(ctx context.Context, idRange ClosedRange) ([]TEntity, error)
+	FetchAllEntities(ctx context.Context) ([]TEntity, error)
 }
 
 type entityFetcherImpl[TEntity Entity] struct {
 	name string
 
 	fetchLastEntityId func(ctx context.Context) (int64, error)
 	fetchEntity       func(ctx context.Context, id int64) (TEntity, error)
-	fetchEntitiesPage func(ctx context.Context, page uint64, limit uint64) ([]TEntity, error)
+
+	fetchEntitiesPage      func(ctx context.Context, page uint64, limit uint64) ([]TEntity, error)
+	fetchEntitiesPageLimit uint64
 
 	logger log.Logger
 }
@@ -30,14 +33,18 @@ func newEntityFetcher[TEntity Entity](
 	fetchLastEntityId func(ctx context.Context) (int64, error),
 	fetchEntity func(ctx context.Context, id int64) (TEntity, error),
 	fetchEntitiesPage func(ctx context.Context, page uint64, limit uint64) ([]TEntity, error),
+	fetchEntitiesPageLimit uint64,
 	logger log.Logger,
 ) entityFetcher[TEntity] {
 	return &entityFetcherImpl[TEntity]{
 		name:              name,
 		fetchLastEntityId: fetchLastEntityId,
 		fetchEntity:       fetchEntity,
-		fetchEntitiesPage: fetchEntitiesPage,
-		logger:            logger,
+
+		fetchEntitiesPage:      fetchEntitiesPage,
+		fetchEntitiesPageLimit: fetchEntitiesPageLimit,
+
+		logger: logger,
 	}
 }
 
@@ -82,7 +89,7 @@ func (f *entityFetcherImpl[TEntity]) FetchAllEntities(ctx context.Context) ([]TE
 	defer progressLogTicker.Stop()
 
 	for page := uint64(1); ; page++ {
-		entitiesPage, err := f.fetchEntitiesPage(ctx, page, 10_000)
+		entitiesPage, err := f.fetchEntitiesPage(ctx, page, f.fetchEntitiesPageLimit)
 		if err != nil {
 			return nil, err
 		}

polygon/heimdall/entity_fetcher_test.go

@@ -0,0 +1,52 @@
+package heimdall
+
+import (
+	"context"
+	"testing"
+
+	"github.com/ledgerwatch/log/v3"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	libcommon "github.com/ledgerwatch/erigon-lib/common"
+)
+
+func TestEntityFetcher_FetchAllEntities(t *testing.T) {
+	for count := uint64(0); count < 20; count++ {
+		testEntityFetcher_FetchAllEntities(t, count, 5)
+	}
+}
+
+func testEntityFetcher_FetchAllEntities(t *testing.T, count uint64, fetchEntitiesPageLimit uint64) {
+	ctx := context.Background()
+	logger := log.New()
+
+	var expectedEntities []*Checkpoint
+	for i := uint64(0); i < count; i++ {
+		c := makeCheckpoint(i*256, 256)
+		c.Id = CheckpointId(i + 1)
+		expectedEntities = append(expectedEntities, c)
+	}
+
+	servedEntities := make([]*Checkpoint, len(expectedEntities))
+	copy(servedEntities, expectedEntities)
+	libcommon.SliceShuffle(servedEntities)
+	fetchEntitiesPage := func(ctx context.Context, page uint64, limit uint64) ([]*Checkpoint, error) {
+		offset := (page - 1) * limit
+		totalLen := uint64(len(servedEntities))
+		return servedEntities[min(offset, totalLen):min(offset+limit, totalLen)], nil
+	}
+
+	fetcher := newEntityFetcher[*Checkpoint](
+		"fetcher",
+		nil,
+		nil,
+		fetchEntitiesPage,
+		fetchEntitiesPageLimit,
+		logger,
+	)
+
+	actualEntities, err := fetcher.FetchAllEntities(ctx)
+	require.NoError(t, err)
+	assert.Equal(t, expectedEntities, actualEntities)
+}

polygon/heimdall/scraper.go

@@ -121,6 +121,7 @@ func newCheckpointFetcher(client HeimdallClient, logger log.Logger) entityFetche
 		client.FetchCheckpointCount,
 		client.FetchCheckpoint,
 		client.FetchCheckpoints,
+		10_000, // fetchEntitiesPageLimit
 		logger,
 	)
 }
@@ -131,6 +132,7 @@ func newMilestoneFetcher(client HeimdallClient, logger log.Logger) entityFetcher
 		client.FetchMilestoneCount,
 		client.FetchMilestone,
 		nil,
+		0,
 		logger,
 	)
 }
@@ -153,6 +155,7 @@ func newSpanFetcher(client HeimdallClient, logger log.Logger) entityFetcher[*Spa
 		fetchLastEntityId,
 		fetchEntity,
 		nil,
+		0,
 		logger,
 	)
 }