From 8334dde12fbc12878f7cf463f13c6ef421c59f6a Mon Sep 17 00:00:00 2001
From: Renan DelValle <rdelval@users.noreply.github.com>
Date: Thu, 28 Sep 2017 16:50:46 -0700
Subject: [PATCH] Sample client now blocks until all hosts entered desired
 state. Cleaned up host maintenance monitor.

---
 examples/client.go | 23 +++++++++++++++++++++++
 monitors.go        | 11 ++++-------
 2 files changed, 27 insertions(+), 7 deletions(-)

diff --git a/examples/client.go b/examples/client.go
index 86972b1..8e173ca 100644
--- a/examples/client.go
+++ b/examples/client.go
@@ -511,6 +511,18 @@ func main() {
 			fmt.Printf("error: %+v\n", err.Error())
 			os.Exit(1)
 		}
+
+		// Monitor change to DRAINING and DRAINED mode
+		_, err = monitor.HostMaintenance(
+			hosts,
+			[]aurora.MaintenanceMode{aurora.MaintenanceMode_DRAINED, aurora.MaintenanceMode_DRAINING},
+			5,
+			10)
+		if err != nil {
+			fmt.Printf("error: %+v\n", err.Error())
+			os.Exit(1)
+		}
+
 		fmt.Print(result.String())
 
 	case "endMaintenance":
@@ -525,6 +537,17 @@ func main() {
 			fmt.Printf("error: %+v\n", err.Error())
 			os.Exit(1)
 		}
+
+		// Monitor change to DRAINING and DRAINED mode
+		_, err = monitor.HostMaintenance(
+			hosts,
+			[]aurora.MaintenanceMode{aurora.MaintenanceMode_NONE},
+			5,
+			10)
+		if err != nil {
+			fmt.Printf("error: %+v\n", err.Error())
+			os.Exit(1)
+		}
 		fmt.Print(result.String())
 
 	default:
diff --git a/monitors.go b/monitors.go
index 69285c0..36caf2e 100644
--- a/monitors.go
+++ b/monitors.go
@@ -161,17 +161,15 @@ func (m *Monitor) HostMaintenance(hosts []string, modes []aurora.MaintenanceMode
 	for _,mode := range modes {
 		desiredMode[mode] = struct{}{}
 	}
-	hostMode := make(map[string]bool)
 
 	// Initial map has all hosts we're looking for.
 	// For each node we find in the correct mode, eliminate it from the map. If we reach 0 elements in the map,
 	// we found all hosts we we're monitoring. This avoids having to go through and check the list one by one each cycle.
+	hostMode := make(map[string]struct{})
 	for _,host := range hosts {
-		hostMode[host] = true
+		hostMode[host] = struct{}{}
 	}
 
-	fmt.Println("mode map and hosts have the same number of elements: ", len(hostMode) == len(hosts))
-
 	for step := 0; step < steps; step++ {
 
 		// Client may have multiple retries handle retries
@@ -183,14 +181,13 @@ func (m *Monitor) HostMaintenance(hosts []string, modes []aurora.MaintenanceMode
 
 		for stat := range result.GetStatuses() {
 			if  _, ok := desiredMode[stat.GetMode()]; ok {
-				fmt.Printf("host %s\n", stat.GetHost())
-				fmt.Println(hostMode)
+				fmt.Printf("host %s entered %s state\n", stat.GetHost(), stat.GetMode())
 				delete(hostMode, stat.GetHost())
 			}
 		}
 
 		if len(hostMode) == 0 {
-			fmt.Println("Provided hosts have all entered the desired state(s)")
+			fmt.Println("Provided hosts have all entered desired state(s)")
 			return true, nil
 		} else {
 			fmt.Printf("%d host(s) not in desired state\n", len(hostMode))