github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/module/epochs/qc_voter_test.go (about)

     1  package epochs_test
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"math/rand"
     7  	"testing"
     8  
     9  	"github.com/rs/zerolog"
    10  	"github.com/stretchr/testify/mock"
    11  	"github.com/stretchr/testify/suite"
    12  
    13  	hotstuff "github.com/onflow/flow-go/consensus/hotstuff/mocks"
    14  	"github.com/onflow/flow-go/model/flow"
    15  	"github.com/onflow/flow-go/model/flow/factory"
    16  	flowmodule "github.com/onflow/flow-go/module"
    17  	"github.com/onflow/flow-go/module/epochs"
    18  	module "github.com/onflow/flow-go/module/mock"
    19  	protocol "github.com/onflow/flow-go/state/protocol/mock"
    20  	"github.com/onflow/flow-go/utils/unittest"
    21  )
    22  
    23  type Suite struct {
    24  	suite.Suite
    25  
    26  	local  *module.Local
    27  	signer *hotstuff.Signer
    28  
    29  	client *module.QCContractClient
    30  	voted  bool
    31  
    32  	state *protocol.State
    33  	snap  *protocol.Snapshot
    34  
    35  	epoch      *protocol.Epoch
    36  	counter    uint64
    37  	phase      flow.EpochPhase
    38  	nodes      flow.IdentityList
    39  	me         *flow.Identity
    40  	clustering flow.ClusterList // cluster assignment for epoch
    41  
    42  	voter *epochs.RootQCVoter
    43  }
    44  
    45  func (suite *Suite) SetupTest() {
    46  
    47  	log := zerolog.New(io.Discard)
    48  	suite.local = new(module.Local)
    49  	suite.signer = new(hotstuff.Signer)
    50  	suite.client = new(module.QCContractClient)
    51  
    52  	suite.voted = false
    53  	suite.client.On("Voted", mock.Anything).Return(
    54  		func(_ context.Context) bool { return suite.voted },
    55  		func(_ context.Context) error { return nil },
    56  	)
    57  	suite.client.On("SubmitVote", mock.Anything, mock.Anything).Return(nil)
    58  
    59  	suite.state = new(protocol.State)
    60  	suite.snap = new(protocol.Snapshot)
    61  	suite.state.On("Final").Return(suite.snap)
    62  	suite.phase = flow.EpochPhaseSetup
    63  	suite.snap.On("Phase").Return(
    64  		func() flow.EpochPhase { return suite.phase },
    65  		func() error { return nil },
    66  	)
    67  
    68  	suite.epoch = new(protocol.Epoch)
    69  	suite.counter = rand.Uint64()
    70  
    71  	suite.nodes = unittest.IdentityListFixture(4, unittest.WithRole(flow.RoleCollection))
    72  	suite.me = suite.nodes[rand.Intn(len(suite.nodes))]
    73  	suite.local.On("NodeID").Return(func() flow.Identifier {
    74  		return suite.me.NodeID
    75  	})
    76  
    77  	var err error
    78  	assignments := unittest.ClusterAssignment(2, suite.nodes.ToSkeleton())
    79  	suite.clustering, err = factory.NewClusterList(assignments, suite.nodes.ToSkeleton())
    80  	suite.Require().NoError(err)
    81  
    82  	suite.epoch.On("Counter").Return(suite.counter, nil)
    83  	suite.epoch.On("Clustering").Return(suite.clustering, nil)
    84  	suite.signer.On("CreateVote", mock.Anything).Return(unittest.VoteFixture(), nil)
    85  
    86  	suite.voter = epochs.NewRootQCVoter(log, suite.local, suite.signer, suite.state, []flowmodule.QCContractClient{suite.client})
    87  }
    88  
    89  func TestRootQCVoter(t *testing.T) {
    90  	suite.Run(t, new(Suite))
    91  }
    92  
    93  // should fail if this node isn't in any cluster next epoch
    94  func (suite *Suite) TestNonClusterParticipant() {
    95  
    96  	// change our identity so we aren't in the cluster assignment
    97  	suite.me = unittest.IdentityFixture(unittest.WithRole(flow.RoleCollection))
    98  	err := suite.voter.Vote(context.Background(), suite.epoch)
    99  	suite.Assert().Error(err)
   100  	suite.Assert().True(epochs.IsClusterQCNoVoteError(err))
   101  }
   102  
   103  // should fail if we are not in setup phase
   104  func (suite *Suite) TestInvalidPhase() {
   105  
   106  	suite.phase = flow.EpochPhaseStaking
   107  	err := suite.voter.Vote(context.Background(), suite.epoch)
   108  	suite.Assert().Error(err)
   109  	suite.Assert().True(epochs.IsClusterQCNoVoteError(err))
   110  }
   111  
   112  // should succeed and exit if we've already voted
   113  func (suite *Suite) TestAlreadyVoted() {
   114  
   115  	suite.voted = true
   116  	err := suite.voter.Vote(context.Background(), suite.epoch)
   117  	suite.Assert().NoError(err)
   118  }
   119  
   120  // should succeed and exit if voting succeeds
   121  func (suite *Suite) TestVoting() {
   122  	err := suite.voter.Vote(context.Background(), suite.epoch)
   123  	suite.Assert().NoError(err)
   124  }