Adding parameter for Aurora so that we're able to run SLA aware updates with less than 20 instances. Lowered time it takes to run test by reducing watch time per instance as well.

realis_e2e_test.go

@@ -750,3 +750,76 @@ func TestRealisClient_PartitionPolicy(t *testing.T) {
 	// Clean up after finishing test
 	_, err = r.KillJob(job.JobKey())
 }
+
+func TestAuroraJob_UpdateSlaPolicy(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"
+
+	_, err := r.SetQuota(role, thrift.Float64Ptr(6.0), thrift.Int64Ptr(1024), thrift.Int64Ptr(1024))
+	assert.NoError(t, err)
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+
+			// Create a single job
+			job := realis.NewJob().
+				Environment("prod").
+				Role(role).
+				Name(tt.name).
+				ExecutorName(aurora.AURORA_EXECUTOR_NAME).
+				ExecutorData(string(thermosPayload)).
+				CPU(.01).
+				RAM(2).
+				Disk(5).
+				InstanceCount(8).
+				IsService(true).
+				SlaPolicy(&tt.args).
+				Tier("preferred")
+
+			settings := realis.NewUpdateSettings()
+			settings.UpdateGroupSize = 4
+			settings.MinWaitInInstanceRunningMs = 5 * 1000
+
+			_, result, err := r.CreateService(job, settings)
+			assert.NoError(t, err)
+			assert.NotNil(t, result)
+
+			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)
+		})
+	}
+}