From a698e3b8dddf0d24b490fad2f9b0980cbdec763b Mon Sep 17 00:00:00 2001
From: Renan DelValle <commit@ridv.xyz>
Date: Thu, 12 Sep 2019 16:57:07 -0700
Subject: [PATCH] Adding TerminalUpdateStates function which returns a slice
 containing all terminal states for an update. Changed signature of
 JobUpdateStatus from using a map for desired states to a slice. A map is no
 longer necessary with the new version of thrift and only adds complexity.

---
 CHANGELOG.md |  5 +++++
 monitors.go  | 29 +++++++----------------------
 realis.go    |  3 +--
 util.go      | 12 ++++++++++++
 4 files changed, 25 insertions(+), 24 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index cfed746..551bae9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,11 @@ Users can check for a Timedout error by using `realis.IsTimeout(err)`.
 a Variable Batch Update configured Update is currently in.
 * Added new PauseUpdateMonitor which monitors an update until it is an `ROLL_FORWARD_PAUSED` state.
 * Added variableBatchStep command to sample client to be used for testing new VariableBatchStep api.
+* JobUpdateStatus has changed function signature from:
+`JobUpdateStatus(updateKey aurora.JobUpdateKey, desiredStatuses map[aurora.JobUpdateStatus]bool, interval, timeout time.Duration) (aurora.JobUpdateStatus, error)`
+to
+`JobUpdateStatus(updateKey aurora.JobUpdateKey, desiredStatuses []aurora.JobUpdateStatus, interval, timeout time.Duration) (aurora.JobUpdateStatus, error)`
+* Added TerminalUpdateStates function which returns an slice containing all UpdateStates which are considered terminal states.
 
 1.21.0
 
diff --git a/monitors.go b/monitors.go
index cd9aa06..452261c 100644
--- a/monitors.go
+++ b/monitors.go
@@ -36,15 +36,9 @@ func (m *Monitor) JobUpdate(
 	timeout int) (bool, error) {
 
 	updateQ := aurora.JobUpdateQuery{
-		Key:   &updateKey,
-		Limit: 1,
-		UpdateStatuses: []aurora.JobUpdateStatus{
-			aurora.JobUpdateStatus_ROLLED_FORWARD,
-			aurora.JobUpdateStatus_ROLLED_BACK,
-			aurora.JobUpdateStatus_ABORTED,
-			aurora.JobUpdateStatus_ERROR,
-			aurora.JobUpdateStatus_FAILED,
-		},
+		Key:            &updateKey,
+		Limit:          1,
+		UpdateStatuses: TerminalUpdateStates(),
 	}
 	updateSummaries, err := m.JobUpdateQuery(
 		updateQ,
@@ -75,22 +69,13 @@ func (m *Monitor) JobUpdate(
 }
 
 // JobUpdateStatus polls the scheduler every certain amount of time to see if the update has entered a specified state.
-func (m *Monitor) JobUpdateStatus(
-	updateKey aurora.JobUpdateKey,
-	desiredStatuses map[aurora.JobUpdateStatus]bool,
-	interval time.Duration,
-	timeout time.Duration) (aurora.JobUpdateStatus, error) {
-
-	desiredStatusesSlice := make([]aurora.JobUpdateStatus, 0)
-
-	for k := range desiredStatuses {
-		desiredStatusesSlice = append(desiredStatusesSlice, k)
-	}
-
+func (m *Monitor) JobUpdateStatus(updateKey aurora.JobUpdateKey,
+	desiredStatuses []aurora.JobUpdateStatus,
+	interval, timeout time.Duration) (aurora.JobUpdateStatus, error) {
 	updateQ := aurora.JobUpdateQuery{
 		Key:            &updateKey,
 		Limit:          1,
-		UpdateStatuses: desiredStatusesSlice,
+		UpdateStatuses: desiredStatuses,
 	}
 	summary, err := m.JobUpdateQuery(updateQ, interval, timeout)
 
diff --git a/realis.go b/realis.go
index bd0efb2..28127b2 100644
--- a/realis.go
+++ b/realis.go
@@ -830,7 +830,7 @@ func (r *realisClient) AbortJobUpdate(updateKey aurora.JobUpdateKey, message str
 	m := Monitor{Client: r}
 	_, err := m.JobUpdateStatus(
 		updateKey,
-		map[aurora.JobUpdateStatus]bool{aurora.JobUpdateStatus_ABORTED: true},
+		[]aurora.JobUpdateStatus{aurora.JobUpdateStatus_ABORTED},
 		time.Second*5,
 		time.Minute)
 
@@ -1063,7 +1063,6 @@ func (r *realisClient) RollbackJobUpdate(key aurora.JobUpdateKey, message string
 
 // VariableBatchStep returns the current batch the update is in during a variable batch update.
 func (r *realisClient) VariableBatchStep(key aurora.JobUpdateKey) (int, error) {
-
 	// Query Aurora for an Update that is either paused or rolling forward
 	resp, err := r.JobUpdateDetails(aurora.JobUpdateQuery{
 		UpdateStatuses: []aurora.JobUpdateStatus{
diff --git a/util.go b/util.go
index 8431347..08f3490 100644
--- a/util.go
+++ b/util.go
@@ -25,6 +25,18 @@ var TerminalStates = make(map[aurora.ScheduleStatus]bool)
 // ActiveJobUpdateStates - States a Job Update may be in where it is considered active.
 var ActiveJobUpdateStates = make(map[aurora.JobUpdateStatus]bool)
 
+// TerminalJobUpdateStates returns a slice containing all the terminal states an update may end up in.
+// This is a function in order to avoid having a slice that can be accidentally mutated.
+func TerminalUpdateStates() []aurora.JobUpdateStatus {
+	return []aurora.JobUpdateStatus{
+		aurora.JobUpdateStatus_ROLLED_FORWARD,
+		aurora.JobUpdateStatus_ROLLED_BACK,
+		aurora.JobUpdateStatus_ABORTED,
+		aurora.JobUpdateStatus_ERROR,
+		aurora.JobUpdateStatus_FAILED,
+	}
+}
+
 // AwaitingPulseJobUpdateStates - States a job update may be in where it is waiting for a pulse.
 var AwaitingPulseJobUpdateStates = make(map[aurora.JobUpdateStatus]bool)