From db9bebb80270daa939d75244d7f1f5ec488abc65 Mon Sep 17 00:00:00 2001
From: "Tan N. Le" <lenhattan86@users.noreply.github.com>
Date: Mon, 1 Nov 2021 18:17:49 -0700
Subject: [PATCH] enable default sla for slaDrain (#138)

---
 realis.go          |  9 +++++++++
 realis_admin.go    | 12 ++++++++++++
 realis_e2e_test.go | 47 ++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 68 insertions(+)

diff --git a/realis.go b/realis.go
index 772e9ab..89de748 100644
--- a/realis.go
+++ b/realis.go
@@ -129,6 +129,15 @@ var defaultBackoff = Backoff{
 	Jitter:   0.1,
 }
 
+var defaultSlaPolicy = aurora.SlaPolicy{
+	PercentageSlaPolicy: &aurora.PercentageSlaPolicy{
+		Percentage:   66,
+		DurationSecs: 300,
+	},
+}
+
+const defaultSlaDrainTimeoutSecs = 900
+
 // ClientOption is an alias for a function that modifies the realis config object
 type ClientOption func(*config)
 
diff --git a/realis_admin.go b/realis_admin.go
index cec92af..9c58081 100644
--- a/realis_admin.go
+++ b/realis_admin.go
@@ -58,6 +58,18 @@ func (r *realisClient) SLADrainHosts(
 		return nil, errors.New("no hosts provided to drain")
 	}
 
+	if policy == nil || policy.CountSetFieldsSlaPolicy() == 0 {
+		policy = &defaultSlaPolicy
+		r.logger.Printf("Warning: start draining with default sla policy %v", policy)
+	}
+
+	if timeout < 0 {
+		r.logger.Printf("Warning: timeout %d secs is invalid, draining with default timeout %d secs",
+			timeout,
+			defaultSlaDrainTimeoutSecs)
+		timeout = defaultSlaDrainTimeoutSecs
+	}
+
 	drainList := aurora.NewHosts()
 	drainList.HostNames = hosts
 
diff --git a/realis_e2e_test.go b/realis_e2e_test.go
index 99155ab..78d8f0e 100644
--- a/realis_e2e_test.go
+++ b/realis_e2e_test.go
@@ -750,6 +750,53 @@ func TestRealisClient_SLADrainHosts(t *testing.T) {
 		5,
 		10)
 	assert.NoError(t, err)
+
+	// slaDrainHosts goes with default policy if no policy is specified
+	_, err = r.SLADrainHosts(nil, 30, hosts...)
+	require.NoError(t, err, "unable to drain host with SLA policy")
+
+	// Monitor change to DRAINING and DRAINED mode
+	hostResults, err = monitor.HostMaintenance(
+		hosts,
+		[]aurora.MaintenanceMode{aurora.MaintenanceMode_DRAINED, aurora.MaintenanceMode_DRAINING},
+		1,
+		50)
+	assert.NoError(t, err)
+	assert.Equal(t, map[string]bool{"localhost": true}, hostResults)
+
+	_, _, err = r.EndMaintenance(hosts...)
+	require.NoError(t, err)
+
+	// Monitor change to DRAINING and DRAINED mode
+	_, err = monitor.HostMaintenance(
+		hosts,
+		[]aurora.MaintenanceMode{aurora.MaintenanceMode_NONE},
+		5,
+		10)
+	assert.NoError(t, err)
+
+	_, err = r.SLADrainHosts(&aurora.SlaPolicy{}, 30, hosts...)
+	require.NoError(t, err, "unable to drain host with SLA policy")
+
+	// Monitor change to DRAINING and DRAINED mode
+	hostResults, err = monitor.HostMaintenance(
+		hosts,
+		[]aurora.MaintenanceMode{aurora.MaintenanceMode_DRAINED, aurora.MaintenanceMode_DRAINING},
+		1,
+		50)
+	assert.NoError(t, err)
+	assert.Equal(t, map[string]bool{"localhost": true}, hostResults)
+
+	_, _, err = r.EndMaintenance(hosts...)
+	require.NoError(t, err)
+
+	// Monitor change to DRAINING and DRAINED mode
+	_, err = monitor.HostMaintenance(
+		hosts,
+		[]aurora.MaintenanceMode{aurora.MaintenanceMode_NONE},
+		5,
+		10)
+	assert.NoError(t, err)
 }
 
 // Test multiple go routines using a single connection