From 9ed7690de14bb8bcc4808a6132d0a5f37bfdcebe Mon Sep 17 00:00:00 2001
From: Robert Allen <zircote@gmail.com>
Date: Wed, 19 Dec 2018 15:24:44 -0600
Subject: [PATCH] Adds support for Tier and SlaPolicy to the Job interface

This adds the following function to the Job interface
- `Tier(tier *string) Job`
- `SlaPolicy(policy *aurora.SlaPolicy) Job`
- Test coverage for each in the end to end tests.
---
 job.go             | 15 ++++++++
 realis_e2e_test.go | 86 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 101 insertions(+)

diff --git a/job.go b/job.go
index eb7eaa0..b1ad23f 100644
--- a/job.go
+++ b/job.go
@@ -56,6 +56,8 @@ type Job interface {
 	MaxFailure(maxFail int32) Job
 	Container(container Container) Job
 	PartitionPolicy(policy *aurora.PartitionPolicy) Job
+	Tier(tier *string) Job
+	SlaPolicy(policy *aurora.SlaPolicy) Job
 }
 
 // Structure to collect all information pertaining to an Aurora job.
@@ -320,6 +322,19 @@ func (j *AuroraJob) Container(container Container) Job {
 // Set a partition policy for the job configuration to implement.
 func (j *AuroraJob) PartitionPolicy(policy *aurora.PartitionPolicy) Job {
 	j.jobConfig.TaskConfig.PartitionPolicy = policy
+	return j
+}
+
+// Set the Tier for the Job.
+func (j *AuroraJob) Tier(tier *string) Job {
+	j.jobConfig.TaskConfig.Tier = tier
+
+	return j
+}
+
+// Set an SlaPolicy for the Job.
+func (j *AuroraJob) SlaPolicy(policy *aurora.SlaPolicy) Job {
+	j.jobConfig.TaskConfig.SlaPolicy = policy
 
 	return j
 }
diff --git a/realis_e2e_test.go b/realis_e2e_test.go
index 7fb3504..5b055a2 100644
--- a/realis_e2e_test.go
+++ b/realis_e2e_test.go
@@ -727,6 +727,92 @@ func TestRealisClient_PartitionPolicy(t *testing.T) {
 		_, err = r.KillJob(job.JobKey())
 
 		assert.NoError(t, err)
+
+	}
+
+}
+
+func TestAuroraJob_SlaPolicy(t *testing.T) {
+
+	tests := []struct {
+		name string
+		args aurora.SlaPolicy
+	}{
+		{
+			"create_service_with_sla_count_policy_test",
+			aurora.SlaPolicy{CountSlaPolicy: &aurora.CountSlaPolicy{Count: 1, DurationSecs: 30}},
+		},
+		{
+			"create_service_with_sla_percentage_policy_test",
+			aurora.SlaPolicy{PercentageSlaPolicy: &aurora.PercentageSlaPolicy{Percentage: 0.25, DurationSecs: 30}},
+		},
+		{
+			"create_service_with_sla_coordinator_policy_test",
+			aurora.SlaPolicy{CoordinatorSlaPolicy: &aurora.CoordinatorSlaPolicy{
+				CoordinatorUrl: "http://localhost/endpoint", StatusKey: "aurora_test"}},
+		},
+	}
+	role := "vagrant"
+	tier := "preferred"
+
+	var cpu float64 = 6.0
+	var disk int64 = 256
+	var ram int64 = 256
+
+	resp, err := r.SetQuota(role, &cpu, &ram, &disk)
+
+	assert.NoError(t, err)
+	assert.Equal(t, aurora.ResponseCode_OK, resp.ResponseCode)
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+
+			// Create a single job
+			job := realis.NewJob().
+				Environment("prod").
+				Role("vagrant").
+				Name(tt.name).
+				ExecutorName(aurora.AURORA_EXECUTOR_NAME).
+				ExecutorData(string(thermosPayload)).
+				CPU(.1).
+				RAM(4).
+				Disk(10).
+				InstanceCount(3).
+				IsService(true).
+				SlaPolicy(&tt.args).
+				Tier(&tier)
+
+			settings := realis.NewUpdateSettings()
+			settings.UpdateGroupSize = 1
+			resp, result, err := r.CreateService(job, settings)
+
+			assert.NoError(t, err)
+			assert.NotNil(t, result)
+			assert.Equal(t, aurora.ResponseCode_OK, resp.GetResponseCode())
+
+			var ok bool
+			var mErr error
+
+			if ok, mErr = monitor.JobUpdate(*result.GetKey(), 5, 240); !ok || mErr != nil {
+				// Update may already be in a terminal state so don't check for error
+				_, err := r.AbortJobUpdate(*result.GetKey(), "Monitor timed out.")
+
+				_, err = r.KillJob(job.JobKey())
+
+				assert.NoError(t, err)
+			}
+			assert.True(t, ok)
+			assert.NoError(t, mErr)
+
+			// Kill task test task after confirming it came up fine
+			_, err = r.KillJob(job.JobKey())
+
+			assert.NoError(t, err)
+
+			cpu = 0
+			ram = 0
+			disk = 0
+			r.SetQuota(role, &cpu, &ram, &disk)
+		})
 	}
 
 }