github.com/mailgun/holster/v4@v4.20.0/election/cluster_test.go (about)

     1  package election_test
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"sync"
     7  	"testing"
     8  
     9  	"github.com/mailgun/holster/v4/election"
    10  	"github.com/mailgun/holster/v4/setter"
    11  	"github.com/stretchr/testify/require"
    12  )
    13  
    14  type ChangePair struct {
    15  	From   string
    16  	Leader string
    17  }
    18  
    19  // Useful in tests where you need to simulate an election cluster
    20  type TestCluster struct {
    21  	Nodes      map[string]*ClusterNode
    22  	OnChangeCh chan ChangePair
    23  	t          *testing.T
    24  	errors     map[string]error
    25  	lock       sync.Mutex
    26  }
    27  
    28  type ClusterNode struct {
    29  	lock    sync.RWMutex
    30  	Node    election.Node
    31  	SendRPC func(from string, to string, req election.RPCRequest, resp *election.RPCResponse) error
    32  }
    33  
    34  func NewTestCluster(t *testing.T) *TestCluster {
    35  	return &TestCluster{
    36  		Nodes:      make(map[string]*ClusterNode),
    37  		t:          t,
    38  		errors:     make(map[string]error),
    39  		OnChangeCh: make(chan ChangePair, 500),
    40  	}
    41  }
    42  
    43  // Spawns a new node and adds it to the cluster
    44  func (c *TestCluster) SpawnNode(name string, conf *election.Config) error {
    45  	setter.SetDefault(&conf, &election.Config{})
    46  	n := &ClusterNode{
    47  		SendRPC: c.sendRPC,
    48  	}
    49  
    50  	conf.UniqueID = name
    51  	conf.SendRPC = func(ctx context.Context, peer string, req election.RPCRequest, resp *election.RPCResponse) error {
    52  		n.lock.RLock()
    53  		defer n.lock.RUnlock()
    54  		return n.SendRPC(name, peer, req, resp)
    55  	}
    56  	conf.OnUpdate = func(s string) {
    57  		c.OnChangeCh <- ChangePair{
    58  			From:   name,
    59  			Leader: s,
    60  		}
    61  	}
    62  	var err error
    63  	n.Node, err = election.NewNode(*conf)
    64  	if err != nil {
    65  		return err
    66  	}
    67  	// Add the node to our list of nodes
    68  	c.Add(name, n)
    69  	err = n.Node.Start(context.Background())
    70  	require.NoError(c.t, err)
    71  	return nil
    72  }
    73  
    74  func (c *TestCluster) Add(name string, node *ClusterNode) {
    75  	c.lock.Lock()
    76  	defer c.lock.Unlock()
    77  	c.Nodes[name] = node
    78  
    79  	node.lock.Lock()
    80  	defer node.lock.Unlock()
    81  	node.SendRPC = c.sendRPC
    82  	c.updatePeers()
    83  }
    84  
    85  func (c *TestCluster) Remove(name string) *ClusterNode {
    86  	c.lock.Lock()
    87  	defer c.lock.Unlock()
    88  
    89  	n := c.Nodes[name]
    90  	delete(c.Nodes, name)
    91  	c.updatePeers()
    92  	return n
    93  }
    94  
    95  func (c *TestCluster) updatePeers() {
    96  	// Build a list of all the peers
    97  	var peers []string
    98  	for k := range c.Nodes {
    99  		peers = append(peers, k)
   100  	}
   101  
   102  	// Update our list of known peers
   103  	for _, v := range c.Nodes {
   104  		err := v.Node.SetPeers(context.Background(), peers)
   105  		require.NoError(c.t, err)
   106  	}
   107  }
   108  
   109  type ClusterStatus map[string]string
   110  
   111  func (c *TestCluster) GetClusterStatus() ClusterStatus {
   112  	status := make(ClusterStatus)
   113  	for k, v := range c.Nodes {
   114  		status[k] = v.Node.GetLeader()
   115  	}
   116  	return status
   117  }
   118  
   119  func (c *TestCluster) GetLeader() election.Node {
   120  	for _, v := range c.Nodes {
   121  		if v.Node.IsLeader() {
   122  			return v.Node
   123  		}
   124  	}
   125  	return nil
   126  }
   127  
   128  func (c *TestCluster) peerKey(from, to string) string {
   129  	return fmt.Sprintf("%s|%s", from, to)
   130  }
   131  
   132  func (c *TestCluster) ClearErrors() {
   133  	c.lock.Lock()
   134  	defer c.lock.Unlock()
   135  	c.errors = make(map[string]error)
   136  }
   137  
   138  // Add a specific peer to peer error
   139  func (c *TestCluster) Disconnect(from, to string, err error) {
   140  	c.lock.Lock()
   141  	defer c.lock.Unlock()
   142  	c.errors[c.peerKey(from, to)] = err
   143  }
   144  
   145  func (c *TestCluster) AddNetworkError(peer string, err error) {
   146  	c.lock.Lock()
   147  	defer c.lock.Unlock()
   148  	c.errors[peer] = err
   149  }
   150  
   151  func (c *TestCluster) DelNetworkError(peer string) {
   152  	c.lock.Lock()
   153  	defer c.lock.Unlock()
   154  	delete(c.errors, peer)
   155  }
   156  
   157  func (c *TestCluster) sendRPC(from, to string, req election.RPCRequest, resp *election.RPCResponse) error {
   158  	c.lock.Lock()
   159  	defer c.lock.Unlock()
   160  
   161  	if err, ok := c.errors[from]; ok {
   162  		return err
   163  	}
   164  
   165  	if err, ok := c.errors[to]; ok {
   166  		return err
   167  	}
   168  
   169  	if err, ok := c.errors[c.peerKey(from, to)]; ok {
   170  		return err
   171  	}
   172  
   173  	n, ok := c.Nodes[to]
   174  	if !ok {
   175  		return fmt.Errorf("unknown peer '%s'", to)
   176  	}
   177  	n.Node.ReceiveRPC(req, resp)
   178  
   179  	return nil
   180  }
   181  
   182  func (c *TestCluster) Close() {
   183  	for _, v := range c.Nodes {
   184  		err := v.Node.Stop(context.Background())
   185  		require.NoError(c.t, err)
   186  	}
   187  }