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