github.com/MetalBlockchain/metalgo@v1.11.9/snow/engine/snowman/block/test_batched_vm.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package block
     5  
     6  import (
     7  	"context"
     8  	"errors"
     9  	"testing"
    10  	"time"
    11  
    12  	"github.com/stretchr/testify/require"
    13  
    14  	"github.com/MetalBlockchain/metalgo/ids"
    15  	"github.com/MetalBlockchain/metalgo/snow/consensus/snowman"
    16  )
    17  
    18  var (
    19  	errGetAncestor       = errors.New("unexpectedly called GetAncestor")
    20  	errBatchedParseBlock = errors.New("unexpectedly called BatchedParseBlock")
    21  
    22  	_ BatchedChainVM = (*TestBatchedVM)(nil)
    23  )
    24  
    25  // TestBatchedVM is a BatchedVM that is useful for testing.
    26  type TestBatchedVM struct {
    27  	T *testing.T
    28  
    29  	CantGetAncestors    bool
    30  	CantBatchParseBlock bool
    31  
    32  	GetAncestorsF func(
    33  		ctx context.Context,
    34  		blkID ids.ID,
    35  		maxBlocksNum int,
    36  		maxBlocksSize int,
    37  		maxBlocksRetrivalTime time.Duration,
    38  	) ([][]byte, error)
    39  
    40  	BatchedParseBlockF func(
    41  		ctx context.Context,
    42  		blks [][]byte,
    43  	) ([]snowman.Block, error)
    44  }
    45  
    46  func (vm *TestBatchedVM) Default(cant bool) {
    47  	vm.CantGetAncestors = cant
    48  	vm.CantBatchParseBlock = cant
    49  }
    50  
    51  func (vm *TestBatchedVM) GetAncestors(
    52  	ctx context.Context,
    53  	blkID ids.ID,
    54  	maxBlocksNum int,
    55  	maxBlocksSize int,
    56  	maxBlocksRetrivalTime time.Duration,
    57  ) ([][]byte, error) {
    58  	if vm.GetAncestorsF != nil {
    59  		return vm.GetAncestorsF(
    60  			ctx,
    61  			blkID,
    62  			maxBlocksNum,
    63  			maxBlocksSize,
    64  			maxBlocksRetrivalTime,
    65  		)
    66  	}
    67  	if vm.CantGetAncestors && vm.T != nil {
    68  		require.FailNow(vm.T, errGetAncestor.Error())
    69  	}
    70  	return nil, errGetAncestor
    71  }
    72  
    73  func (vm *TestBatchedVM) BatchedParseBlock(
    74  	ctx context.Context,
    75  	blks [][]byte,
    76  ) ([]snowman.Block, error) {
    77  	if vm.BatchedParseBlockF != nil {
    78  		return vm.BatchedParseBlockF(ctx, blks)
    79  	}
    80  	if vm.CantBatchParseBlock && vm.T != nil {
    81  		require.FailNow(vm.T, errBatchedParseBlock.Error())
    82  	}
    83  	return nil, errBatchedParseBlock
    84  }