github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/utils/unittest/headers.go (about)

     1  package unittest
     2  
     3  import (
     4  	"github.com/stretchr/testify/mock"
     5  
     6  	"github.com/onflow/flow-go/model/flow"
     7  	storerr "github.com/onflow/flow-go/storage"
     8  	storage "github.com/onflow/flow-go/storage/mock"
     9  )
    10  
    11  // HeadersFromMap creates a storage header mock that backed by a given map
    12  func HeadersFromMap(headerDB map[flow.Identifier]*flow.Header) *storage.Headers {
    13  	headers := &storage.Headers{}
    14  	headers.On("Store", mock.Anything).Return(
    15  		func(header *flow.Header) error {
    16  			headerDB[header.ID()] = header
    17  			return nil
    18  		},
    19  	)
    20  	headers.On("ByBlockID", mock.Anything).Return(
    21  		func(blockID flow.Identifier) *flow.Header {
    22  			return headerDB[blockID]
    23  		},
    24  		func(blockID flow.Identifier) error {
    25  			_, exists := headerDB[blockID]
    26  			if !exists {
    27  				return storerr.ErrNotFound
    28  			}
    29  			return nil
    30  		},
    31  	)
    32  
    33  	return headers
    34  }