github.com/annchain/OG@v0.0.9/consensus/dkg/dkg_partner_test.go (about)

     1  package dkg
     2  
     3  import (
     4  	"github.com/annchain/kyber/v3/pairing/bn256"
     5  	"github.com/sirupsen/logrus"
     6  	"sync"
     7  	"testing"
     8  	"time"
     9  )
    10  
    11  var TestNodes = 4
    12  
    13  func init() {
    14  	Formatter := new(logrus.TextFormatter)
    15  	//Formatter.ForceColors = false
    16  	Formatter.DisableColors = true
    17  	Formatter.TimestampFormat = "15:04:05.000000"
    18  	Formatter.FullTimestamp = true
    19  	logrus.SetLevel(logrus.InfoLevel)
    20  	logrus.SetFormatter(Formatter)
    21  	//logrus.SetReportCaller(true)
    22  
    23  	//filenameHook := filename.NewHook()
    24  	//filenameHook.Field = "line"
    25  	//logrus.AddHook(filenameHook)
    26  }
    27  
    28  func setupPartners(termId uint32, numParts int, threshold int) (dkgPartners []*DefaultDkgPartner, partSecs []PartSec, err error) { // generate dkger
    29  	suite := bn256.NewSuiteG2()
    30  
    31  	dkgers, partSecs, err := SetupAllDkgers(suite, numParts, threshold)
    32  	if err != nil {
    33  		return
    34  	}
    35  
    36  	// setup partners
    37  	peerChans := make([]chan *DkgMessageEvent, numParts)
    38  
    39  	// prepare incoming channels
    40  	for i := 0; i < numParts; i++ {
    41  		peerChans[i] = make(chan *DkgMessageEvent, 5000)
    42  	}
    43  
    44  	dkgPartners = make([]*DefaultDkgPartner, numParts)
    45  	partPubs := make([]PartPub, numParts)
    46  
    47  	for i, partSec := range partSecs {
    48  		partPubs[i] = partSec.PartPub
    49  	}
    50  
    51  	for i, dkger := range dkgers {
    52  		communicator := NewDummyDkgPeerCommunicator(i, peerChans[i], peerChans)
    53  		communicator.Run()
    54  
    55  		partner, err := NewDefaultDkgPartner(suite, termId, numParts, threshold, partPubs, partSecs[i],
    56  			communicator, communicator)
    57  		if err != nil {
    58  			panic(err)
    59  		}
    60  		partner.context.Dkger = dkger
    61  		dkgPartners[i] = partner
    62  	}
    63  	return
    64  
    65  }
    66  
    67  type dummyDkgGeneratedListener struct {
    68  	Wg       *sync.WaitGroup
    69  	c        chan bool
    70  	finished bool
    71  }
    72  
    73  func NewDummyDkgGeneratedListener(wg *sync.WaitGroup) *dummyDkgGeneratedListener {
    74  	d := &dummyDkgGeneratedListener{
    75  		Wg: wg,
    76  		c:  make(chan bool),
    77  	}
    78  	go func() {
    79  		for {
    80  			<-d.c
    81  			d.Wg.Done()
    82  			logrus.Info("Dkg is generated")
    83  			//break
    84  		}
    85  	}()
    86  	return d
    87  }
    88  
    89  func (d dummyDkgGeneratedListener) GetDkgGeneratedEventChannel() chan bool {
    90  	return d.c
    91  }
    92  
    93  func TestDkgPartner(t *testing.T) {
    94  	// simulate 4 dkg partners
    95  	termId := uint32(0)
    96  	numParts := TestNodes
    97  	threshold := TestNodes
    98  
    99  	dkgPartners, partSecs, err := setupPartners(termId, numParts, threshold)
   100  	if err != nil {
   101  		panic(err)
   102  	}
   103  
   104  	wg := sync.WaitGroup{}
   105  
   106  	wg.Add(len(partSecs)) // no need to use partSecs, just for note
   107  	listener := NewDummyDkgGeneratedListener(&wg)
   108  
   109  	for _, partner := range dkgPartners {
   110  		partner.dkgGeneratedListeners = append(partner.dkgGeneratedListeners, listener)
   111  		partner.Start()
   112  	}
   113  	time.Sleep(time.Second * 5)
   114  	for _, partner := range dkgPartners {
   115  		partner.gossipStartCh <- true
   116  	}
   117  
   118  	wg.Wait()
   119  }