github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/consensus/hotstuff/votecollector/testutil.go (about)

     1  package votecollector
     2  
     3  import (
     4  	"github.com/onflow/crypto"
     5  	"github.com/stretchr/testify/mock"
     6  	"github.com/stretchr/testify/suite"
     7  
     8  	"github.com/onflow/flow-go/consensus/hotstuff/helper"
     9  	mockhotstuff "github.com/onflow/flow-go/consensus/hotstuff/mocks"
    10  	"github.com/onflow/flow-go/consensus/hotstuff/model"
    11  	"github.com/onflow/flow-go/model/flow"
    12  )
    13  
    14  // VoteProcessorTestSuiteBase is a helper structure which implements common logic between staking and combined vote
    15  // processor test suites.
    16  type VoteProcessorTestSuiteBase struct {
    17  	suite.Suite
    18  
    19  	sigWeight          uint64
    20  	stakingTotalWeight uint64
    21  	onQCCreatedState   mock.Mock
    22  
    23  	stakingAggregator *mockhotstuff.WeightedSignatureAggregator
    24  	minRequiredWeight uint64
    25  	proposal          *model.Proposal
    26  }
    27  
    28  func (s *VoteProcessorTestSuiteBase) SetupTest() {
    29  	s.stakingAggregator = &mockhotstuff.WeightedSignatureAggregator{}
    30  	s.proposal = helper.MakeProposal()
    31  
    32  	// let's assume we have 19 nodes each with weight 100
    33  	s.sigWeight = 100
    34  	s.minRequiredWeight = 1300 // we require at least 13 sigs to collect min weight
    35  	s.stakingTotalWeight = 0
    36  
    37  	// setup staking signature aggregator
    38  	s.stakingAggregator.On("TrustedAdd", mock.Anything, mock.Anything).Run(func(args mock.Arguments) {
    39  		s.stakingTotalWeight += s.sigWeight
    40  	}).Return(func(signerID flow.Identifier, sig crypto.Signature) uint64 {
    41  		return s.stakingTotalWeight
    42  	}, func(signerID flow.Identifier, sig crypto.Signature) error {
    43  		return nil
    44  	}).Maybe()
    45  	s.stakingAggregator.On("TotalWeight").Return(func() uint64 {
    46  		return s.stakingTotalWeight
    47  	}).Maybe()
    48  }
    49  
    50  // onQCCreated is a special function that registers call in mocked state.
    51  // ATTENTION: don't change name of this function since the same name is used in:
    52  // s.onQCCreatedState.On("onQCCreated") statements
    53  func (s *VoteProcessorTestSuiteBase) onQCCreated(qc *flow.QuorumCertificate) {
    54  	s.onQCCreatedState.Called(qc)
    55  }