add unit test

realis_e2e_test.go

@@ -871,3 +871,59 @@ func TestRealisClient_GetJobSummary(t *testing.T) {
 
 	assert.NoError(t, err)
 }
+
+func TestRealisClient_SlaRestartInstances(t *testing.T) {
+	// Create a single job
+	role := "vagrant"
+	env := "prod"
+	name := "slaRestartInstances"
+
+	job := realis.NewJob().
+		Environment(env).
+		Role(role).
+		Name(name).
+		ThermosExecutor(thermosExec).
+		CPU(.01).
+		RAM(4).
+		Disk(10).
+		Tier("preferred").
+		InstanceCount(3).
+		IsService(true)
+
+	// Needed to populate the task config correctly
+	assert.NoError(t, job.BuildThermosPayload())
+
+	var cpu = 3.5
+	var ram int64 = 20480
+	var disk int64 = 10240
+	err := r.SetQuota(role, &cpu, &ram, &disk)
+	assert.NoError(t, err)
+
+	err = r.CreateJob(job)
+	assert.NoError(t, err)
+
+	// waiting until all instances running
+	success, err := r.MonitorScheduleStatus(job.JobKey(),
+		job.GetInstanceCount(),
+		[]aurora.ScheduleStatus{aurora.ScheduleStatus_RUNNING},
+		1*time.Second,
+		150*time.Second)
+	assert.True(t, success)
+	assert.NoError(t, err)
+
+	slaPolicy := &aurora.SlaPolicy{
+		PercentageSlaPolicy: &aurora.PercentageSlaPolicy{
+			Percentage:   50,
+			DurationSecs: 0,
+		},
+	}
+
+	t.Run("TestRealisClient_SlaRestartInstances", func(t *testing.T) {
+		result, err := r.SlaRestartInstances(job.JobKey(), slaPolicy, 0)
+		assert.NoError(t, err)
+		assert.NotNil(t, result)
+
+		assert.NoError(t, r.AbortJobUpdate(*result.GetKey(), "abort update to kill the job"))
+		assert.NoError(t, r.KillJob(job.JobKey()))
+	})
+}