github.com/niedbalski/juju@v0.0.0-20190215020005-8ff100488e47/worker/raft/rafttest/fixtures.go (about)

     1  // Copyright 2018 Canonical Ltd.
     2  // Licensed under the AGPLv3, see LICENCE file for details.
     3  
     4  package rafttest
     5  
     6  import (
     7  	"log"
     8  	"time"
     9  
    10  	"github.com/hashicorp/raft"
    11  	"github.com/juju/loggo"
    12  	"github.com/juju/testing"
    13  	jc "github.com/juju/testing/checkers"
    14  	gc "gopkg.in/check.v1"
    15  
    16  	coretesting "github.com/juju/juju/testing"
    17  	"github.com/juju/juju/worker/raft/raftutil"
    18  )
    19  
    20  // RaftFixture is a fixture to embed into test suites, providing
    21  // a raft.Raft and in-memory transport.
    22  type RaftFixture struct {
    23  	testing.IsolationSuite
    24  
    25  	FSM           raft.FSM
    26  	Config        *raft.Config
    27  	Transport     *raft.InmemTransport
    28  	Store         *raft.InmemStore
    29  	SnapshotStore *raft.InmemSnapshotStore
    30  	Raft          *raft.Raft
    31  }
    32  
    33  func (s *RaftFixture) SetUpTest(c *gc.C) {
    34  	s.IsolationSuite.SetUpTest(c)
    35  	c.Assert(s.FSM, gc.NotNil, gc.Commentf("FSM must be set by embedding test suite"))
    36  
    37  	s.Raft, s.Config, s.Transport, s.Store, s.SnapshotStore = s.NewRaft(c, "0", s.FSM)
    38  	c.Assert(s.Raft.BootstrapCluster(raft.Configuration{
    39  		Servers: []raft.Server{{
    40  			ID:      s.Config.LocalID,
    41  			Address: s.Transport.LocalAddr(),
    42  		}},
    43  	}).Error(), jc.ErrorIsNil)
    44  	select {
    45  	case leader := <-s.Raft.LeaderCh():
    46  		c.Assert(leader, jc.IsTrue)
    47  	case <-time.After(coretesting.LongWait):
    48  		c.Fatal("timed out waiting for raft leadership")
    49  	}
    50  }
    51  
    52  func (s *RaftFixture) NewRaft(c *gc.C, id raft.ServerID, fsm raft.FSM) (
    53  	*raft.Raft,
    54  	*raft.Config,
    55  	*raft.InmemTransport,
    56  	*raft.InmemStore,
    57  	*raft.InmemSnapshotStore,
    58  ) {
    59  	_, transport := raft.NewInmemTransport("")
    60  	s.AddCleanup(func(*gc.C) { transport.Close() })
    61  
    62  	store := raft.NewInmemStore()
    63  	snapshotStore := raft.NewInmemSnapshotStore()
    64  	config := s.DefaultConfig(id)
    65  
    66  	r, err := raft.NewRaft(config, fsm, store, store, snapshotStore, transport)
    67  	c.Assert(err, jc.ErrorIsNil)
    68  	s.AddCleanup(func(c *gc.C) {
    69  		f := r.Shutdown()
    70  		c.Assert(f.Error(), jc.ErrorIsNil)
    71  	})
    72  	return r, config, transport, store, snapshotStore
    73  }
    74  
    75  func (s *RaftFixture) DefaultConfig(id raft.ServerID) *raft.Config {
    76  	raftConfig := raft.DefaultConfig()
    77  	raftConfig.ShutdownOnRemove = false
    78  	raftConfig.LocalID = id
    79  	raftConfig.HeartbeatTimeout = 100 * time.Millisecond
    80  	raftConfig.ElectionTimeout = raftConfig.HeartbeatTimeout
    81  	raftConfig.LeaderLeaseTimeout = raftConfig.HeartbeatTimeout
    82  	raftConfig.Logger = log.New(&raftutil.LoggoWriter{
    83  		loggo.GetLogger("juju.worker.raft.raftclusterer_test_" + string(id)),
    84  		loggo.DEBUG,
    85  	}, "", 0)
    86  	return raftConfig
    87  }