Test validator's ability to rewind upon span rotation

consensus/bor/bor.go

@@ -1724,6 +1724,17 @@ func (c *Bor) SetHeimdallClient(h IHeimdallClient) {
 	c.spanStore.setHeimdallClient(h)
 }
 
+// PurgeCache clears all cached snapshots and span data. This is useful in tests
+// when the mock heimdall client is changed and old cached data needs to be invalidated.
+func (c *Bor) PurgeCache() {
+	// Clear the recents cache (snapshots)
+	c.recents.DeleteAll()
+	// Clear the recent verified headers cache
+	c.recentVerifiedHeaders.DeleteAll()
+	// Clear the span store cache
+	c.spanStore.PurgeCache()
+}
+
 func (c *Bor) GetCurrentValidators(ctx context.Context, headerHash common.Hash, blockNumber uint64) ([]*valset.Validator, error) {
 	return c.spanner.GetCurrentValidatorsByHash(ctx, headerHash, blockNumber)
 }

consensus/bor/bor_test.go

@@ -1313,3 +1313,50 @@ func TestFinalizeAndAssembleReturnsCommitTime(t *testing.T) {
 		require.Contains(t, processingErr.Error(), "failed to decode genesis alloc")
 	})
 }
+
+func TestBor_PurgeCache(t *testing.T) {
+	t.Parallel()
+	borConfig := &params.BorConfig{
+		Period:                map[string]uint64{"0": 2},
+		ProducerDelay:         map[string]uint64{"0": 4},
+		Sprint:                map[string]uint64{"0": 64},
+		BackupMultiplier:      map[string]uint64{"0": 2},
+		ValidatorContract:     "0x0000000000000000000000000000000000001000",
+		StateReceiverContract: "0x0000000000000000000000000000000000001001",
+	}
+	accountAddr := common.HexToAddress("0x1234567890123456789012345678901234567890")
+	sp := &fakeSpanner{vals: []*valset.Validator{{Address: accountAddr, VotingPower: 50}}}
+	_, borObj := newChainAndBorForTest(t, sp, borConfig, true, accountAddr, uint64(time.Now().Unix()))
+
+	// Add some entries to the recents cache (snapshots)
+	hash1 := common.HexToHash("0x1111111111111111111111111111111111111111111111111111111111111111")
+	hash2 := common.HexToHash("0x2222222222222222222222222222222222222222222222222222222222222222")
+
+	snapshot1 := &Snapshot{Number: 1, Hash: hash1}
+	snapshot2 := &Snapshot{Number: 2, Hash: hash2}
+
+	borObj.recents.Set(hash1, snapshot1, ttlcache.DefaultTTL)
+	borObj.recents.Set(hash2, snapshot2, ttlcache.DefaultTTL)
+
+	// Add some entries to the recentVerifiedHeaders cache
+	header1 := &types.Header{Number: big.NewInt(1)}
+	header2 := &types.Header{Number: big.NewInt(2)}
+
+	borObj.recentVerifiedHeaders.Set(hash1, header1, ttlcache.DefaultTTL)
+	borObj.recentVerifiedHeaders.Set(hash2, header2, ttlcache.DefaultTTL)
+
+	// Verify caches are populated
+	require.Equal(t, 2, borObj.recents.Len(), "recents cache should have 2 entries")
+	require.Equal(t, 2, borObj.recentVerifiedHeaders.Len(), "recentVerifiedHeaders cache should have 2 entries")
+
+	// Purge the cache
+	borObj.PurgeCache()
+
+	// Verify caches are cleared
+	require.Equal(t, 0, borObj.recents.Len(), "recents cache should be empty after purge")
+	require.Equal(t, 0, borObj.recentVerifiedHeaders.Len(), "recentVerifiedHeaders cache should be empty after purge")
+
+	// Verify we can still add entries after purge
+	borObj.recents.Set(hash1, snapshot1, ttlcache.DefaultTTL)
+	require.Equal(t, 1, borObj.recents.Len(), "should be able to add to recents cache after purge")
+}

consensus/bor/span_store.go

@@ -359,6 +359,20 @@ func (s *SpanStore) setHeimdallClient(client IHeimdallClient) {
 	s.heimdallClient = client
 }
 
+// PurgeCache clears all cached spans and resets state. This is useful in tests
+// when the mock heimdall client is changed and old cached data needs to be invalidated.
+func (s *SpanStore) PurgeCache() {
+	// Create a new cache to replace the old one
+	newCache, _ := lru.NewARC(10)
+	s.store = newCache
+	// Clear the latest span cache
+	s.latestSpanCache.Store(nil)
+	// Clear the last used span
+	s.lastUsedSpan.Store(nil)
+	// Reset the latest known span ID
+	s.latestKnownSpanId.Store(0)
+}
+
 // getMockSpan0 constructs a mock span 0 by fetching validator set from genesis state. This should
 // only be used in tests where heimdall client is not available.
 func getMockSpan0(ctx context.Context, spanner Spanner, chainId string) (*borTypes.Span, error) {

consensus/bor/span_store_test.go

@@ -1526,3 +1526,54 @@ func TestSpanStore_HeimdallDownTimeout(t *testing.T) {
 		require.False(t, status.CatchingUp, "Status should show not catching up after recovery")
 	})
 }
+
+func TestSpanStore_PurgeCache(t *testing.T) {
+	t.Parallel()
+	spanStore := NewSpanStore(&MockHeimdallClient{}, nil, "1337")
+	defer spanStore.Close()
+	ctx := t.Context()
+
+	// Populate the cache with some spans
+	for i := uint64(0); i < 5; i++ {
+		span, err := spanStore.spanById(ctx, i)
+		require.NoError(t, err, "err in spanById for id=%d", i)
+		require.NotNil(t, span)
+	}
+
+	// Verify cache is populated
+	keys := spanStore.store.Keys()
+	require.Len(t, keys, 5, "cache should have 5 spans")
+
+	// Set up lastUsedSpan and latestKnownSpanId
+	span, err := spanStore.spanById(ctx, 2)
+	require.NoError(t, err)
+	spanStore.lastUsedSpan.Store(span)
+	spanStore.latestKnownSpanId.Store(4)
+	spanStore.latestSpanCache.Store(span)
+
+	// Verify values are set
+	require.NotNil(t, spanStore.lastUsedSpan.Load(), "lastUsedSpan should be set")
+	require.Equal(t, uint64(4), spanStore.latestKnownSpanId.Load(), "latestKnownSpanId should be 4")
+	require.NotNil(t, spanStore.latestSpanCache.Load(), "latestSpanCache should be set")
+
+	// Purge the cache
+	spanStore.PurgeCache()
+
+	// Verify cache is cleared
+	keys = spanStore.store.Keys()
+	require.Len(t, keys, 0, "cache should be empty after purge")
+
+	// Verify atomic values are reset
+	require.Nil(t, spanStore.lastUsedSpan.Load(), "lastUsedSpan should be nil after purge")
+	require.Equal(t, uint64(0), spanStore.latestKnownSpanId.Load(), "latestKnownSpanId should be 0 after purge")
+	require.Nil(t, spanStore.latestSpanCache.Load(), "latestSpanCache should be nil after purge")
+
+	// Verify we can still fetch spans after purge (cache miss should re-fetch)
+	span, err = spanStore.spanById(ctx, 0)
+	require.NoError(t, err, "should be able to fetch span after purge")
+	require.Equal(t, uint64(0), span.Id, "span id should be 0")
+
+	// Verify cache is being populated again
+	keys = spanStore.store.Keys()
+	require.Len(t, keys, 1, "cache should have 1 span after re-fetch")
+}

eth/ethconfig/config.go

@@ -209,6 +209,10 @@ type Config struct {
 	// Use child heimdall process to fetch data, Only works when RunHeimdall is true
 	UseHeimdallApp bool
 
+	// OverrideHeimdallClient allows injecting a mock HeimdallClient for testing.
+	// When set, this client is used instead of creating one from HeimdallURL/HeimdallgRPCAddress.
+	OverrideHeimdallClient bor.IHeimdallClient `toml:"-"`
+
 	// Bor logs flag
 	BorLogs bool
 
@@ -285,7 +289,10 @@ func CreateConsensusEngine(chainConfig *params.ChainConfig, ethConfig *Config, d
 			}
 
 			var heimdallClient bor.IHeimdallClient
-			if ethConfig.RunHeimdall && ethConfig.UseHeimdallApp {
+			// Use override client if provided (for testing)
+			if ethConfig.OverrideHeimdallClient != nil {
+				heimdallClient = ethConfig.OverrideHeimdallClient
+			} else if ethConfig.RunHeimdall && ethConfig.UseHeimdallApp {
 				// TODO: Running heimdall from bor is not tested yet.
 				// heimdallClient = heimdallapp.NewHeimdallAppClient()
 				panic("Running heimdall from bor is not implemented yet. Please use heimdall gRPC or HTTP client instead.")

eth/ethconfig/config_test.go

@@ -0,0 +1,101 @@
+package ethconfig
+
+import (
+	"context"
+	"math/big"
+	"testing"
+
+	borTypes "github.com/0xPolygon/heimdall-v2/x/bor/types"
+	stakeTypes "github.com/0xPolygon/heimdall-v2/x/stake/types"
+	ctypes "github.com/cometbft/cometbft/rpc/core/types"
+	"github.com/ethereum/go-ethereum/consensus/bor"
+	"github.com/ethereum/go-ethereum/consensus/bor/clerk"
+	"github.com/ethereum/go-ethereum/consensus/bor/heimdall/checkpoint"
+	"github.com/ethereum/go-ethereum/consensus/bor/heimdall/milestone"
+	"github.com/ethereum/go-ethereum/core/rawdb"
+	"github.com/ethereum/go-ethereum/params"
+	"github.com/stretchr/testify/require"
+)
+
+// mockHeimdallClient implements bor.IHeimdallClient for testing
+type mockHeimdallClient struct{}
+
+func (m *mockHeimdallClient) Close() {}
+func (m *mockHeimdallClient) StateSyncEvents(context.Context, uint64, int64) ([]*clerk.EventRecordWithTime, error) {
+	return nil, nil
+}
+func (m *mockHeimdallClient) GetSpan(_ context.Context, spanID uint64) (*borTypes.Span, error) {
+	return &borTypes.Span{
+		Id: spanID, StartBlock: 0, EndBlock: 255,
+		ValidatorSet: stakeTypes.ValidatorSet{
+			Validators: []*stakeTypes.Validator{{ValId: 1, Signer: "0x96C42C56fdb78294F96B0cFa33c92bed7D75F96a", VotingPower: 100}},
+		},
+	}, nil
+}
+func (m *mockHeimdallClient) GetLatestSpan(ctx context.Context) (*borTypes.Span, error) {
+	return m.GetSpan(ctx, 0)
+}
+func (m *mockHeimdallClient) FetchCheckpoint(context.Context, int64) (*checkpoint.Checkpoint, error) {
+	return nil, nil
+}
+func (m *mockHeimdallClient) FetchCheckpointCount(context.Context) (int64, error) { return 0, nil }
+func (m *mockHeimdallClient) FetchMilestone(context.Context) (*milestone.Milestone, error) {
+	return nil, nil
+}
+func (m *mockHeimdallClient) FetchMilestoneCount(context.Context) (int64, error) { return 0, nil }
+func (m *mockHeimdallClient) FetchStatus(context.Context) (*ctypes.SyncInfo, error) {
+	return &ctypes.SyncInfo{CatchingUp: false}, nil
+}
+
+// newTestBorChainConfig creates a minimal Bor chain config for testing
+func newTestBorChainConfig() *params.ChainConfig {
+	return &params.ChainConfig{
+		ChainID:             big.NewInt(137),
+		HomesteadBlock:      big.NewInt(0),
+		EIP150Block:         big.NewInt(0),
+		EIP155Block:         big.NewInt(0),
+		EIP158Block:         big.NewInt(0),
+		ByzantiumBlock:      big.NewInt(0),
+		ConstantinopleBlock: big.NewInt(0),
+		PetersburgBlock:     big.NewInt(0),
+		IstanbulBlock:       big.NewInt(0),
+		MuirGlacierBlock:    big.NewInt(0),
+		BerlinBlock:         big.NewInt(0),
+		LondonBlock:         big.NewInt(0),
+		Bor: &params.BorConfig{
+			Period:                map[string]uint64{"0": 2},
+			ProducerDelay:         map[string]uint64{"0": 4},
+			Sprint:                map[string]uint64{"0": 64},
+			BackupMultiplier:      map[string]uint64{"0": 2},
+			ValidatorContract:     "0x0000000000000000000000000000000000001000",
+			StateReceiverContract: "0x0000000000000000000000000000000000001001",
+		},
+	}
+}
+
+func TestCreateConsensusEngine_OverrideHeimdallClient(t *testing.T) {
+	t.Parallel()
+	ethConfig := &Config{
+		OverrideHeimdallClient: &mockHeimdallClient{},
+		WithoutHeimdall:        false,
+	}
+
+	engine, err := CreateConsensusEngine(newTestBorChainConfig(), ethConfig, rawdb.NewMemoryDatabase(), nil)
+	require.NoError(t, err)
+	defer engine.Close()
+
+	_, ok := engine.(*bor.Bor)
+	require.True(t, ok, "Expected Bor consensus engine")
+}
+
+func TestCreateConsensusEngine_WithoutHeimdall(t *testing.T) {
+	t.Parallel()
+	ethConfig := &Config{WithoutHeimdall: true}
+
+	engine, err := CreateConsensusEngine(newTestBorChainConfig(), ethConfig, rawdb.NewMemoryDatabase(), nil)
+	require.NoError(t, err)
+	defer engine.Close()
+
+	_, ok := engine.(*bor.Bor)
+	require.True(t, ok, "Expected Bor consensus engine")
+}