unittest for election.go

consensus/consensus_test.go

@@ -0,0 +1 @@
+package consensus

consensus/election/election.go

@@ -43,11 +43,11 @@ type Election struct {
 // NewElection creates an election using the config provided.
 func NewElection(config *Config) (*Election, error) {
 	if config.Duration == 0 {
-		return nil, errors.New("Election duration cannot be empty")
+		return nil, errors.New("election duration cannot be empty")
 	}
 
 	if config.MinVotingInterval > config.MaxVotingInterval {
-		return nil, fmt.Errorf("Min voting interval %v is greater than max voting interval %v", config.MinVotingInterval, config.MaxVotingInterval)
+		return nil, fmt.Errorf("min voting interval %v is greater than max voting interval %v", config.MinVotingInterval, config.MaxVotingInterval)
 	}
 
 	if config.GetWeight == nil {
@@ -68,7 +68,7 @@ func NewElection(config *Config) (*Election, error) {
 // returns error.
 func (election *Election) SetInitialVote(vote interface{}) error {
 	if election.HasStarted() {
-		return errors.New("Cannot set initial vote, election has started")
+		return errors.New("cannot set initial vote, election has started")
 	}
 
 	election.Lock()
@@ -100,7 +100,7 @@ func (election *Election) Start() bool {
 	return success
 }
 
-// Stop stops an election. Typically this should not be called directly.
+// Stop stops an election. Typically, this should not be called directly.
 func (election *Election) Stop() {
 	election.Lock()
 	election.state = stopped
@@ -129,7 +129,7 @@ func (election *Election) IsStopped() bool {
 // ReceiveVote receives and saves a vote from a neighbor.
 func (election *Election) ReceiveVote(neighborID, vote interface{}) error {
 	if election.IsStopped() {
-		return errors.New("Election has already stopped")
+		return errors.New("election has already stopped")
 	}
 
 	election.neighborVotes.Store(neighborID, vote)
@@ -142,7 +142,7 @@ func (election *Election) ReceiveVote(neighborID, vote interface{}) error {
 	return nil
 }
 
-// GetTxVoteChan returns the send vote channel, which should be used to send
+// GetTxVoteChan returns the sending vote channel, which should be used to send
 // votes to neighbors.
 func (election *Election) GetTxVoteChan() <-chan interface{} {
 	return election.txVoteChan
@@ -194,7 +194,7 @@ func (election *Election) NeighborVoteCount() uint32 {
 // PrefillNeighborVotes prefills vote for neighborIDs that don't have a vote yet
 func (election *Election) PrefillNeighborVotes(neighborIDs []interface{}, vote interface{}) error {
 	if election.IsStopped() {
-		return errors.New("Election has already stopped")
+		return errors.New("election has already stopped")
 	}
 
 	for _, neighborID := range neighborIDs {

consensus/election/election_test.go

@@ -0,0 +1,229 @@
+package election
+
+import (
+	"encoding/hex"
+	"github.com/nknorg/nkn/v2/common"
+	"github.com/nknorg/nkn/v2/config"
+	"github.com/nknorg/nkn/v2/util"
+	"math/rand"
+	"os"
+	"reflect"
+	"sync"
+	"testing"
+	"time"
+)
+
+var neighbors []*Election
+
+func TestMain(m *testing.M) {
+	neighbors = make([]*Election, 20)
+	for i := 0; i < len(neighbors); i++ {
+		elc, _ := newElection()
+		neighbors[i] = elc
+	}
+	os.Exit(m.Run())
+}
+
+func TestOneVoteElection(t *testing.T) {
+	check := func(f string, got, want interface{}) {
+		if !reflect.DeepEqual(got, want) {
+			t.Errorf("%s mismatch: got %v, want %v", f, got, want)
+		}
+	}
+	elc, err := newElection()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	elc.Start()
+
+	neighborID := "123"
+	err = elc.ReceiveVote(neighborID, common.EmptyUint256)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	elc.Stop()
+	result, u, f, err := elc.GetResult()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	check("result", result, common.EmptyUint256)
+	check("abs weight", u, uint32(1))
+	check("relative weight", f, float32(1))
+}
+
+func TestElection(t *testing.T) {
+	check := func(f string, got, want interface{}) {
+		if !reflect.DeepEqual(got, want) {
+			t.Errorf("%s mismatch: got %v, want %v", f, got, want)
+		}
+	}
+	elc, err := newElection()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	err = elc.SetInitialVote(common.EmptyUint256)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	elc.Start()
+
+	neighborIDs := make([]interface{}, 10)
+	for i := range neighborIDs {
+		neighborIDs[i] = hex.EncodeToString(util.RandomBytes(16))
+	}
+	err = elc.PrefillNeighborVotes(neighborIDs, common.MaxUint256)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	err = elc.ReceiveVote("123", common.EmptyUint256)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	nids := elc.GetNeighborIDsByVote(common.MaxUint256)
+	check("neighbors id", len(nids), len(neighborIDs))
+
+	nid := elc.GetNeighborIDsByVote(common.EmptyUint256)
+	check("neighbor id", len(nid), 1)
+
+	count := elc.NeighborVoteCount()
+	check("neighbor vote count", count, uint32(11))
+
+	elc.Stop()
+	result, u, f, err := elc.GetResult()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	check("result", result, common.MaxUint256)
+	check("abs weight", u, uint32(10))
+	check("relative weight", f, float32(0.8333333))
+}
+
+func TestMultiElection(t *testing.T) {
+	check := func(f string, got, want interface{}) {
+		if !reflect.DeepEqual(got, want) {
+			t.Errorf("%s mismatch: got %v, want %v", f, got, want)
+		}
+	}
+	var wg sync.WaitGroup
+	for i := 0; i < len(neighbors); i++ {
+		i := i
+		wg.Add(1)
+		go func(int) {
+			defer wg.Done()
+			neighbors[i].Start()
+		}(i)
+	}
+	for _, e := range neighbors {
+		go e.run()
+	}
+
+	h := common.EmptyUint256
+	h[0] = 1
+
+	neighbors[0].txVoteChan <- h
+	neighbors[1].txVoteChan <- h
+	neighbors[2].txVoteChan <- common.MaxUint256
+
+	time.Sleep(1 * time.Second)
+
+	wg.Wait()
+	for _, e := range neighbors {
+		for {
+			select {
+			case <-e.txVoteChan:
+				continue
+			default:
+			}
+			break
+		}
+		e.Stop()
+		result, _, _, err := e.GetResult()
+		if err != nil {
+			t.Fatal(err)
+		}
+		check("result", result, h)
+	}
+
+}
+
+func newElection() (*Election, error) {
+	getWeight := func(neighborID interface{}) uint32 {
+		return 1
+	}
+	c := &Config{
+		Duration:                    config.ConsensusDuration / 2,
+		MinVotingInterval:           200 * time.Millisecond,
+		MaxVotingInterval:           2 * time.Second,
+		ChangeVoteMinRelativeWeight: 0.5,
+		ConsensusMinRelativeWeight:  2.0 / 3.0,
+		GetWeight:                   getWeight,
+	}
+
+	elc, err := NewElection(c)
+	if err != nil {
+		return nil, err
+	}
+	return elc, nil
+}
+
+func getRandomElemFromSlice(n []string) string {
+	rand.NewSource(time.Now().Unix())
+	rand.Shuffle(len(n), func(i, j int) {
+		n[i], n[j] = n[j], n[i]
+	})
+	return n[0]
+}
+
+func (election *Election) run() {
+	n1 := []string{"1", "2", "3", "6", "7", "8", "9", "10"}
+	n2 := []string{"4", "5", "11"}
+
+	for {
+		i := 0
+		election.neighborVotes.Range(func(key, value any) bool {
+			i++
+			return true
+		})
+		if i == len(n1)+len(n2) {
+			return
+		}
+		select {
+		case val := <-election.txVoteChan:
+			hash, ok := val.(common.Uint256)
+			if !ok {
+				return
+			}
+
+			var neighborID string
+			if hash == common.MaxUint256 {
+				neighborID = getRandomElemFromSlice(n2)
+			} else {
+				neighborID = getRandomElemFromSlice(n1)
+			}
+			err := election.ReceiveVote(neighborID, hash)
+			if err != nil {
+				return
+			}
+
+			for _, n := range neighbors {
+				nn := n
+				go func() {
+					if nn == election {
+						return
+					}
+					if !nn.IsStopped() {
+						nn.txVoteChan <- hash
+					}
+				}()
+			}
+		}
+	}
+}