From bce64886204bf871794e48ec6ae0449282942d4f Mon Sep 17 00:00:00 2001
From: Ian Gulliver <github@flamingcow.io>
Date: Sun, 4 Jun 2023 10:43:32 -0700
Subject: [PATCH] Add TestSystem

---
 elect_test.go | 21 +++-----------------
 lib_test.go   | 54 +++++++++++++++++++++++++++++++++++++++++++++++++++
 voter.go      | 13 ++++---------
 3 files changed, 61 insertions(+), 27 deletions(-)

diff --git a/elect_test.go b/elect_test.go
index c2607c3..5fc3b9a 100644
--- a/elect_test.go
+++ b/elect_test.go
@@ -1,32 +1,17 @@
 package elect_test
 
 import (
-	"fmt"
 	"testing"
 	"time"
 
-	"github.com/dchest/uniuri"
-	"github.com/gopatchy/elect"
-	"github.com/gopatchy/proxy"
-	"github.com/samber/lo"
 	"github.com/stretchr/testify/require"
 )
 
-func TestSimple(t *testing.T) {
+func TestOne(t *testing.T) {
 	t.Parallel()
 
-	signingKey := uniuri.New()
-
-	ts := NewTestServer(t, signingKey)
+	ts := NewTestSystem(t, 1)
 	defer ts.Stop()
 
-	p := lo.Must(proxy.NewProxy(t, ts.Addr()))
-	defer p.Close()
-
-	url := fmt.Sprintf("http://%s/", p.Addr())
-
-	v := elect.NewVoter(url, signingKey)
-	defer v.Stop()
-
-	require.Eventually(t, ts.Candidate.IsLeader, 15*time.Second, 100*time.Millisecond)
+	require.Eventually(t, ts.Candidate(0).IsLeader, 15*time.Second, 100*time.Millisecond)
 }
diff --git a/lib_test.go b/lib_test.go
index 1e5b77c..7704b30 100644
--- a/lib_test.go
+++ b/lib_test.go
@@ -1,12 +1,15 @@
 package elect_test
 
 import (
+	"fmt"
 	"net"
 	"net/http"
 	"testing"
 	"time"
 
+	"github.com/dchest/uniuri"
 	"github.com/gopatchy/elect"
+	"github.com/gopatchy/proxy"
 	"github.com/samber/lo"
 	"github.com/stretchr/testify/require"
 )
@@ -18,6 +21,13 @@ type TestServer struct {
 	srv      *http.Server
 }
 
+type TestSystem struct {
+	signingKey string
+	servers    []*TestServer
+	voters     []*elect.Voter
+	proxy      *proxy.Proxy
+}
+
 func NewTestServer(t *testing.T, signingKey string) *TestServer {
 	ts := &TestServer{
 		Candidate: elect.NewCandidate(1, signingKey),
@@ -45,3 +55,47 @@ func (ts *TestServer) Stop() {
 func (ts *TestServer) Addr() *net.TCPAddr {
 	return ts.listener.Addr().(*net.TCPAddr)
 }
+
+func NewTestSystem(t *testing.T, num int) *TestSystem {
+	ts := &TestSystem{
+		signingKey: uniuri.New(),
+	}
+
+	for i := 0; i < num; i++ {
+		ts.servers = append(ts.servers, NewTestServer(t, ts.signingKey))
+	}
+
+	ts.proxy = lo.Must(proxy.NewProxy(t, ts.Server(0).Addr()))
+
+	url := fmt.Sprintf("http://%s/", ts.proxy.Addr())
+
+	for i := 0; i < num; i++ {
+		ts.voters = append(ts.voters, elect.NewVoter(url, ts.signingKey, ts.Candidate(i)))
+	}
+
+	return ts
+}
+
+func (ts *TestSystem) Stop() {
+	for _, s := range ts.servers {
+		s.Stop()
+	}
+
+	for _, v := range ts.voters {
+		v.Stop()
+	}
+
+	ts.proxy.Close()
+}
+
+func (ts *TestSystem) Server(i int) *TestServer {
+	return ts.servers[i]
+}
+
+func (ts *TestSystem) Candidate(i int) *elect.Candidate {
+	return ts.servers[i].Candidate
+}
+
+func (ts *TestSystem) Voter(i int) *elect.Voter {
+	return ts.voters[i]
+}
diff --git a/voter.go b/voter.go
index 6b1bc1d..a6c6886 100644
--- a/voter.go
+++ b/voter.go
@@ -20,8 +20,8 @@ type Voter struct {
 	// used by loop() goroutine only
 	client     *resty.Client
 	signingKey []byte
+	candidate  *Candidate
 	vote       vote
-	candidates []*Candidate
 	period     time.Duration
 }
 
@@ -40,12 +40,13 @@ type voteResponse struct {
 	ResponseSent time.Time `json:"responseSent"`
 }
 
-func NewVoter(url string, signingKey string) *Voter {
+func NewVoter(url string, signingKey string, candidate *Candidate) *Voter {
 	v := &Voter{
 		client: resty.New().
 			SetCloseConnection(true).
 			SetBaseURL(url),
 		signingKey: []byte(signingKey),
+		candidate:  candidate,
 		update:     make(chan time.Duration),
 		done:       make(chan bool),
 		vote: vote{
@@ -64,10 +65,6 @@ func (v *Voter) Stop() {
 	<-v.done
 }
 
-func (v *Voter) AddCandidate(c *Candidate) {
-	v.candidates = append(v.candidates, c)
-}
-
 func (v *Voter) loop() {
 	defer close(v.done)
 
@@ -111,9 +108,7 @@ func (v *Voter) poll() bool {
 func (v *Voter) sendVote() {
 	v.vote.VoteSent = time.Now().UTC()
 
-	for _, c := range v.candidates {
-		c.voteIfNo(&v.vote)
-	}
+	v.candidate.voteIfNo(&v.vote)
 
 	js := lo.Must(json.Marshal(v.vote))