From 296af622d1217e28f9ee8573d540c2fc7f836627 Mon Sep 17 00:00:00 2001
From: Robert Allen <zircote@gmail.com>
Date: Thu, 20 Dec 2018 16:38:06 -0600
Subject: [PATCH] This adds the following function to the PartitionPolicy
 configuration to the Job interface (#91)

* Adding Partition Policy API
---
 job.go             |  8 ++++++++
 realis_e2e_test.go | 38 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 46 insertions(+)

diff --git a/job.go b/job.go
index 0ff2aac..eb7eaa0 100644
--- a/job.go
+++ b/job.go
@@ -55,6 +55,7 @@ type Job interface {
 	GetInstanceCount() int32
 	MaxFailure(maxFail int32) Job
 	Container(container Container) Job
+	PartitionPolicy(policy *aurora.PartitionPolicy) Job
 }
 
 // Structure to collect all information pertaining to an Aurora job.
@@ -315,3 +316,10 @@ func (j *AuroraJob) Container(container Container) Job {
 
 	return j
 }
+
+// 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
+}
diff --git a/realis_e2e_test.go b/realis_e2e_test.go
index 688a797..7fb3504 100644
--- a/realis_e2e_test.go
+++ b/realis_e2e_test.go
@@ -692,3 +692,41 @@ func TestRealisClient_ForceExplicitTaskReconciliation(t *testing.T) {
 	err = r.ForceExplicitTaskReconciliation(&batchSize)
 	assert.NoError(t, err)
 }
+
+func TestRealisClient_PartitionPolicy(t *testing.T) {
+
+	role := "vagrant"
+	var partitionDelay int64 = 30
+	job := realis.NewJob().
+		Environment("prod").
+		Role(role).
+		Name("create_thermos_job_partition_policy_test").
+		ExecutorName(aurora.AURORA_EXECUTOR_NAME).
+		ExecutorData(string(thermosPayload)).
+		CPU(.5).
+		RAM(64).
+		Disk(100).
+		IsService(true).
+		InstanceCount(2).
+		PartitionPolicy(&aurora.PartitionPolicy{Reschedule: true, DelaySecs: &partitionDelay})
+
+	settings := realis.NewUpdateSettings()
+	settings.UpdateGroupSize = 2
+	resp, result, err := r.CreateService(job, settings)
+	assert.NoError(t, err)
+
+	assert.Equal(t, aurora.ResponseCode_OK, resp.ResponseCode)
+
+	var ok bool
+	var mErr error
+
+	if ok, mErr = monitor.JobUpdate(*result.GetKey(), 5, 180); !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)
+	}
+
+}