From 8dc2a9c6e1262e89146e13ae7a57aae2cb5113f9 Mon Sep 17 00:00:00 2001
From: James Rasell <jrasell@hashicorp.com>
Date: Thu, 25 Mar 2021 08:52:59 +0100
Subject: [PATCH] api: add Allocation client and server terminal status funcs.

---
 api/allocations.go                            | 22 ++++++
 api/allocations_test.go                       | 74 +++++++++++++++++++
 .../hashicorp/nomad/api/allocations.go        | 22 ++++++
 3 files changed, 118 insertions(+)

diff --git a/api/allocations.go b/api/allocations.go
index 8933e048a..be4058977 100644
--- a/api/allocations.go
+++ b/api/allocations.go
@@ -419,6 +419,28 @@ func (a *Allocation) Stub() *AllocationListStub {
 	}
 }
 
+// ServerTerminalStatus returns true if the desired state of the allocation is
+// terminal.
+func (a *Allocation) ServerTerminalStatus() bool {
+	switch a.DesiredStatus {
+	case AllocDesiredStatusStop, AllocDesiredStatusEvict:
+		return true
+	default:
+		return false
+	}
+}
+
+// ClientTerminalStatus returns true if the client status is terminal and will
+// therefore no longer transition.
+func (a *Allocation) ClientTerminalStatus() bool {
+	switch a.ClientStatus {
+	case AllocClientStatusComplete, AllocClientStatusFailed, AllocClientStatusLost:
+		return true
+	default:
+		return false
+	}
+}
+
 // AllocationListStub is used to return a subset of an allocation
 // during list operations.
 type AllocationListStub struct {
diff --git a/api/allocations_test.go b/api/allocations_test.go
index 0e514a01e..cc306ec88 100644
--- a/api/allocations_test.go
+++ b/api/allocations_test.go
@@ -316,6 +316,80 @@ func TestAllocations_ExecErrors(t *testing.T) {
 	require.Equal(t, err.Error(), fmt.Sprintf("Unknown allocation \"%s\"", allocID))
 }
 
+func TestAllocation_ServerTerminalStatus(t *testing.T) {
+	t.Parallel()
+
+	testCases := []struct {
+		inputAllocation *Allocation
+		expectedOutput  bool
+		name            string
+	}{
+		{
+			inputAllocation: &Allocation{DesiredStatus: AllocDesiredStatusEvict},
+			expectedOutput:  true,
+			name:            "alloc desired status evict",
+		},
+		{
+			inputAllocation: &Allocation{DesiredStatus: AllocDesiredStatusStop},
+			expectedOutput:  true,
+			name:            "alloc desired status stop",
+		},
+		{
+			inputAllocation: &Allocation{DesiredStatus: AllocDesiredStatusRun},
+			expectedOutput:  false,
+			name:            "alloc desired status run",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			require.Equal(t, tc.expectedOutput, tc.inputAllocation.ServerTerminalStatus(), tc.name)
+		})
+	}
+}
+
+func TestAllocation_ClientTerminalStatus(t *testing.T) {
+	t.Parallel()
+
+	testCases := []struct {
+		inputAllocation *Allocation
+		expectedOutput  bool
+		name            string
+	}{
+		{
+			inputAllocation: &Allocation{ClientStatus: AllocClientStatusLost},
+			expectedOutput:  true,
+			name:            "alloc client status lost",
+		},
+		{
+			inputAllocation: &Allocation{ClientStatus: AllocClientStatusFailed},
+			expectedOutput:  true,
+			name:            "alloc client status failed",
+		},
+		{
+			inputAllocation: &Allocation{ClientStatus: AllocClientStatusComplete},
+			expectedOutput:  true,
+			name:            "alloc client status complete",
+		},
+		{
+			inputAllocation: &Allocation{ClientStatus: AllocClientStatusRunning},
+			expectedOutput:  false,
+			name:            "alloc client status complete",
+		},
+		{
+			inputAllocation: &Allocation{ClientStatus: AllocClientStatusPending},
+			expectedOutput:  false,
+			name:            "alloc client status running",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			require.Equal(t, tc.expectedOutput, tc.inputAllocation.ClientTerminalStatus(), tc.name)
+		})
+	}
+}
+
 func TestAllocations_ShouldMigrate(t *testing.T) {
 	t.Parallel()
 	require.True(t, DesiredTransition{Migrate: boolToPtr(true)}.ShouldMigrate())
diff --git a/vendor/github.com/hashicorp/nomad/api/allocations.go b/vendor/github.com/hashicorp/nomad/api/allocations.go
index 8933e048a..be4058977 100644
--- a/vendor/github.com/hashicorp/nomad/api/allocations.go
+++ b/vendor/github.com/hashicorp/nomad/api/allocations.go
@@ -419,6 +419,28 @@ func (a *Allocation) Stub() *AllocationListStub {
 	}
 }
 
+// ServerTerminalStatus returns true if the desired state of the allocation is
+// terminal.
+func (a *Allocation) ServerTerminalStatus() bool {
+	switch a.DesiredStatus {
+	case AllocDesiredStatusStop, AllocDesiredStatusEvict:
+		return true
+	default:
+		return false
+	}
+}
+
+// ClientTerminalStatus returns true if the client status is terminal and will
+// therefore no longer transition.
+func (a *Allocation) ClientTerminalStatus() bool {
+	switch a.ClientStatus {
+	case AllocClientStatusComplete, AllocClientStatusFailed, AllocClientStatusLost:
+		return true
+	default:
+		return false
+	}
+}
+
 // AllocationListStub is used to return a subset of an allocation
 // during list operations.
 type AllocationListStub struct {