github.com/MetalBlockchain/metalgo@v1.11.9/snow/networking/benchlist/test_benchable.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package benchlist
     5  
     6  import (
     7  	"testing"
     8  
     9  	"github.com/stretchr/testify/require"
    10  
    11  	"github.com/MetalBlockchain/metalgo/ids"
    12  )
    13  
    14  type TestBenchable struct {
    15  	T *testing.T
    16  
    17  	CantBenched, CantUnbenched bool
    18  	BenchedF, UnbenchedF       func(chainID ids.ID, validatorID ids.NodeID)
    19  }
    20  
    21  // Default set the default callable value to [cant]
    22  func (b *TestBenchable) Default(cant bool) {
    23  	b.CantBenched = cant
    24  	b.CantUnbenched = cant
    25  }
    26  
    27  func (b *TestBenchable) Benched(chainID ids.ID, validatorID ids.NodeID) {
    28  	if b.BenchedF != nil {
    29  		b.BenchedF(chainID, validatorID)
    30  	} else if b.CantBenched && b.T != nil {
    31  		require.FailNow(b.T, "Unexpectedly called Benched")
    32  	}
    33  }
    34  
    35  func (b *TestBenchable) Unbenched(chainID ids.ID, validatorID ids.NodeID) {
    36  	if b.UnbenchedF != nil {
    37  		b.UnbenchedF(chainID, validatorID)
    38  	} else if b.CantUnbenched && b.T != nil {
    39  		require.FailNow(b.T, "Unexpectedly called Unbenched")
    40  	}
    41  }