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

     1  package mocks
     2  
     3  import "github.com/onflow/flow-go/storage"
     4  
     5  // StorageMapGetter implements a simple generic getter function for mock storage methods.
     6  // This is useful to avoid duplicating boilerplate code for mock storage methods.
     7  //
     8  // Example:
     9  // Instead of the following code:
    10  //
    11  //	results.On("ByID", mock.AnythingOfType("flow.Identifier")).Return(
    12  //		func(resultID flow.Identifier) (*flow.ExecutionResult, error) {
    13  //			if result, ok := s.resultMap[resultID]; ok {
    14  //				return result, nil
    15  //			}
    16  //			return nil, storage.ErrNotFound
    17  //		},
    18  //	)
    19  //
    20  // Use this:
    21  //
    22  //	results.On("ByID", mock.AnythingOfType("flow.Identifier")).Return(
    23  //		mocks.StorageMapGetter(s.resultMap),
    24  //	)
    25  func StorageMapGetter[K comparable, V any](m map[K]V) func(key K) (V, error) {
    26  	return func(key K) (V, error) {
    27  		if val, ok := m[key]; ok {
    28  			return val, nil
    29  		}
    30  		return *new(V), storage.ErrNotFound
    31  	}
    32  }
    33  
    34  // ConvertStorageOutput maps the output type from a getter function to a different type.
    35  // This is useful to avoid maintaining multiple maps for the same data.
    36  //
    37  // Example usage:
    38  //
    39  //	blockMap := map[uint64]*flow.Block{}
    40  //
    41  //	headers.On("BlockIDByHeight", mock.AnythingOfType("uint64")).Return(
    42  //		mocks.ConvertStorageOutput(
    43  //			mocks.StorageMapGetter(s.blockMap),
    44  //			func(block *flow.Block) flow.Identifier { return block.ID() },
    45  //		),
    46  //	)
    47  func ConvertStorageOutput[K comparable, V any, R any](fn func(key K) (V, error), mapper func(V) R) func(key K) (R, error) {
    48  	return func(key K) (R, error) {
    49  		v, err := fn(key)
    50  		if err != nil {
    51  			return *new(R), err
    52  		}
    53  		return mapper(v), err
    54  	}
    55  }