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  }