From ddea4c91e467eadfa63844b26eba85aa5abee121 Mon Sep 17 00:00:00 2001
From: Ian Gulliver <github@flamingcow.io>
Date: Sun, 4 Jun 2023 14:04:15 -0700
Subject: [PATCH] Test partial voting

---
 elect_test.go | 34 ++++++++++++++++++++++++++++++++++
 voter.go      |  5 +++++
 2 files changed, 39 insertions(+)

diff --git a/elect_test.go b/elect_test.go
index 859a98e..4fdd5cc 100644
--- a/elect_test.go
+++ b/elect_test.go
@@ -13,6 +13,8 @@ func TestOne(t *testing.T) {
 	ts := NewTestSystem(t, 1)
 	defer ts.Stop()
 
+	require.False(t, ts.Candidate(0).IsLeader())
+
 	require.Eventually(t, ts.Candidate(0).IsLeader, 20*time.Second, 100*time.Millisecond)
 }
 
@@ -22,7 +24,13 @@ func TestThree(t *testing.T) {
 	ts := NewTestSystem(t, 3)
 	defer ts.Stop()
 
+	require.False(t, ts.Candidate(0).IsLeader())
+	require.False(t, ts.Candidate(1).IsLeader())
+	require.False(t, ts.Candidate(2).IsLeader())
+
 	require.Eventually(t, ts.Candidate(0).IsLeader, 20*time.Second, 100*time.Millisecond)
+	require.False(t, ts.Candidate(1).IsLeader())
+	require.False(t, ts.Candidate(2).IsLeader())
 }
 
 func TestFailover(t *testing.T) {
@@ -31,7 +39,13 @@ func TestFailover(t *testing.T) {
 	ts := NewTestSystem(t, 3)
 	defer ts.Stop()
 
+	require.False(t, ts.Candidate(0).IsLeader())
+	require.False(t, ts.Candidate(1).IsLeader())
+	require.False(t, ts.Candidate(2).IsLeader())
+
 	require.Eventually(t, ts.Candidate(0).IsLeader, 20*time.Second, 100*time.Millisecond)
+	require.False(t, ts.Candidate(1).IsLeader())
+	require.False(t, ts.Candidate(2).IsLeader())
 
 	ts.SetServer(1)
 
@@ -39,6 +53,26 @@ func TestFailover(t *testing.T) {
 
 	// New candidate must not get leadership before old candidate loses it
 	require.False(t, ts.Candidate(1).IsLeader())
+	require.False(t, ts.Candidate(2).IsLeader())
 
 	require.Eventually(t, ts.Candidate(1).IsLeader, 20*time.Second, 100*time.Millisecond)
+	require.False(t, ts.Candidate(0).IsLeader())
+	require.False(t, ts.Candidate(2).IsLeader())
+}
+
+func TestPartialVotes(t *testing.T) {
+	t.Parallel()
+
+	ts := NewTestSystem(t, 3)
+	defer ts.Stop()
+
+	ts.Voter(0).Stop()
+
+	require.False(t, ts.Candidate(0).IsLeader())
+	require.False(t, ts.Candidate(1).IsLeader())
+	require.False(t, ts.Candidate(2).IsLeader())
+
+	require.Eventually(t, ts.Candidate(0).IsLeader, 20*time.Second, 100*time.Millisecond)
+	require.False(t, ts.Candidate(1).IsLeader())
+	require.False(t, ts.Candidate(2).IsLeader())
 }
diff --git a/voter.go b/voter.go
index 5019a7a..bc967f1 100644
--- a/voter.go
+++ b/voter.go
@@ -61,8 +61,13 @@ func NewVoter(url string, signingKey string, candidate *Candidate) *Voter {
 }
 
 func (v *Voter) Stop() {
+	if v.update == nil {
+		return
+	}
+
 	close(v.update)
 	<-v.done
+	v.update = nil
 }
 
 func (v *Voter) loop() {