From 5d0998647a80fd1461563881928833b819c7b429 Mon Sep 17 00:00:00 2001
From: "Tan N. Le" <lenhattan86@users.noreply.github.com>
Date: Mon, 1 Nov 2021 18:15:51 -0700
Subject: [PATCH] default policy for slaDrainHosts (#17)

---
 realis_admin.go    | 12 ++++++++++++
 realis_config.go   | 10 ++++++++++
 realis_e2e_test.go | 33 +++++++++++++++++++++++++++++++++
 3 files changed, 55 insertions(+)

diff --git a/realis_admin.go b/realis_admin.go
index 9530031..f2759ea 100644
--- a/realis_admin.go
+++ b/realis_admin.go
@@ -59,6 +59,18 @@ func (c *Client) SLADrainHosts(policy *aurora.SlaPolicy, timeout int64, hosts ..
 		return nil, errors.New("no hosts provided to drain")
 	}
 
+	if policy == nil || policy.CountSetFieldsSlaPolicy() == 0 {
+		policy = &defaultSlaPolicy
+		c.logger.Printf("Warning: start draining with default sla policy %v", policy)
+	}
+
+	if timeout < 0 {
+		c.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_config.go b/realis_config.go
index 7bb09fc..4b2b96f 100644
--- a/realis_config.go
+++ b/realis_config.go
@@ -19,6 +19,7 @@ import (
 	"time"
 
 	"github.com/apache/thrift/lib/go/thrift"
+	"github.com/aurora-scheduler/gorealis/v2/gen-go/apache/aurora"
 )
 
 type clientConfig struct {
@@ -48,6 +49,15 @@ var defaultBackoff = Backoff{
 	Jitter:   0.1,
 }
 
+var defaultSlaPolicy = aurora.SlaPolicy{
+	PercentageSlaPolicy: &aurora.PercentageSlaPolicy{
+		Percentage:   66,
+		DurationSecs: 300,
+	},
+}
+
+const defaultSlaDrainTimeoutSecs = 900
+
 type TransportProtocol int
 
 const (
diff --git a/realis_e2e_test.go b/realis_e2e_test.go
index 4707a18..55fab19 100644
--- a/realis_e2e_test.go
+++ b/realis_e2e_test.go
@@ -612,6 +612,39 @@ func TestRealisClient_SLADrainHosts(t *testing.T) {
 		5*time.Second,
 		10*time.Second)
 	assert.NoError(t, err)
+
+	// slaDrainHosts goes with default policy if no policy is specified
+	_, err = r.SLADrainHosts(nil, 30, hosts...)
+	if err != nil {
+		fmt.Printf("error: %+v\n", err.Error())
+		os.Exit(1)
+	}
+	hostResults, err = r.MonitorHostMaintenance(
+		hosts,
+		[]aurora.MaintenanceMode{aurora.MaintenanceMode_DRAINED, aurora.MaintenanceMode_DRAINING},
+		1*time.Second,
+		50*time.Second)
+	assert.Equal(t, map[string]bool{"localhost": true}, hostResults)
+	assert.NoError(t, err)
+
+	_, err = r.EndMaintenance(hosts...)
+	assert.NoError(t, err)
+
+	_, err = r.SLADrainHosts(&aurora.SlaPolicy{}, 30, hosts...)
+	if err != nil {
+		fmt.Printf("error: %+v\n", err.Error())
+		os.Exit(1)
+	}
+	hostResults, err = r.MonitorHostMaintenance(
+		hosts,
+		[]aurora.MaintenanceMode{aurora.MaintenanceMode_DRAINED, aurora.MaintenanceMode_DRAINING},
+		1*time.Second,
+		50*time.Second)
+	assert.Equal(t, map[string]bool{"localhost": true}, hostResults)
+	assert.NoError(t, err)
+
+	_, err = r.EndMaintenance(hosts...)
+	assert.NoError(t, err)
 }
 
 // Test multiple go routines using a single connection