From cee8ad6d5f855998441986c5af5c57a4df1fc635 Mon Sep 17 00:00:00 2001
From: Ian Gulliver <github@flamingcow.io>
Date: Sun, 11 Jun 2023 19:14:34 -0700
Subject: [PATCH] Add --elect-force-state

---
 candidate.go | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/candidate.go b/candidate.go
index b24db46..04a6fe3 100644
--- a/candidate.go
+++ b/candidate.go
@@ -3,6 +3,7 @@ package elect
 import (
 	"crypto/hmac"
 	"encoding/json"
+	"flag"
 	"fmt"
 	"io"
 	"log"
@@ -48,6 +49,8 @@ var StateName = map[CandidateState]string{
 	StateNotLeader: "NOT_LEADER",
 }
 
+var electForceState = flag.String("elect-force-state", "", "'', 'leader', 'notleader'")
+
 func NewCandidate(numVoters int, signingKey string) *Candidate {
 	change := make(chan CandidateState, 100)
 
@@ -207,6 +210,32 @@ func (c *Candidate) elect(v *vote) {
 		c.c <- state
 	}()
 
+	switch *electForceState {
+	case "":
+		// Not forced
+
+	case "leader":
+		log.Printf("[elect] state forced to leader")
+
+		state = StateLeader
+
+		return
+
+	case "not-leader":
+		fallthrough
+	case "not_leader":
+		fallthrough
+	case "notleader":
+		log.Printf("[elect] state forced to not leader")
+
+		state = StateNotLeader
+
+		return
+
+	default:
+		panic("invalid --elect-force-state")
+	}
+
 	if v != nil {
 		v.received = time.Now()
 		c.votes[v.VoterID] = v