github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/engine/execution/testutil/finalized_reader.go (about)

     1  package testutil
     2  
     3  import (
     4  	"fmt"
     5  
     6  	"go.uber.org/atomic"
     7  
     8  	"github.com/onflow/flow-go/model/flow"
     9  	"github.com/onflow/flow-go/storage"
    10  	"github.com/onflow/flow-go/utils/unittest"
    11  )
    12  
    13  type MockFinalizedReader struct {
    14  	headerByHeight  map[uint64]*flow.Header
    15  	blockByHeight   map[uint64]*flow.Block
    16  	lowest          uint64
    17  	highest         uint64
    18  	finalizedHeight *atomic.Uint64
    19  	finalizedCalled *atomic.Int64
    20  }
    21  
    22  func NewMockFinalizedReader(initHeight uint64, count int) (*MockFinalizedReader, map[uint64]*flow.Header, uint64) {
    23  	root := unittest.BlockHeaderFixture(unittest.WithHeaderHeight(initHeight))
    24  	blocks := unittest.ChainFixtureFrom(count, root)
    25  	headerByHeight := make(map[uint64]*flow.Header, len(blocks)+1)
    26  	headerByHeight[root.Height] = root
    27  
    28  	blockByHeight := make(map[uint64]*flow.Block, len(blocks)+1)
    29  	for _, b := range blocks {
    30  		headerByHeight[b.Header.Height] = b.Header
    31  		blockByHeight[b.Header.Height] = b
    32  	}
    33  
    34  	highest := blocks[len(blocks)-1].Header.Height
    35  	return &MockFinalizedReader{
    36  		headerByHeight:  headerByHeight,
    37  		blockByHeight:   blockByHeight,
    38  		lowest:          initHeight,
    39  		highest:         highest,
    40  		finalizedHeight: atomic.NewUint64(initHeight),
    41  		finalizedCalled: atomic.NewInt64(0),
    42  	}, headerByHeight, highest
    43  }
    44  
    45  func (r *MockFinalizedReader) FinalizedBlockIDAtHeight(height uint64) (flow.Identifier, error) {
    46  	r.finalizedCalled.Add(1)
    47  	finalized := r.finalizedHeight.Load()
    48  	if height > finalized {
    49  		return flow.Identifier{}, storage.ErrNotFound
    50  	}
    51  
    52  	if height < r.lowest {
    53  		return flow.ZeroID, fmt.Errorf("height %d is out of range [%d, %d]", height, r.lowest, r.highest)
    54  	}
    55  	return r.headerByHeight[height].ID(), nil
    56  }
    57  
    58  func (r *MockFinalizedReader) MockFinal(height uint64) error {
    59  	if height < r.lowest || height > r.highest {
    60  		return fmt.Errorf("height %d is out of range [%d, %d]", height, r.lowest, r.highest)
    61  	}
    62  
    63  	r.finalizedHeight.Store(height)
    64  	return nil
    65  }
    66  
    67  func (r *MockFinalizedReader) BlockAtHeight(height uint64) *flow.Block {
    68  	return r.blockByHeight[height]
    69  }
    70  
    71  func (r *MockFinalizedReader) FinalizedCalled() int {
    72  	return int(r.finalizedCalled.Load())
    73  }