From 089c4396c47f76d5f274734c74a9b6bdf6508769 Mon Sep 17 00:00:00 2001
From: Ryan Uber <ru@ryanuber.com>
Date: Wed, 21 Jan 2015 13:02:47 -0800
Subject: [PATCH] api: support reason flag for maintenance mode

---
 api/agent.go      | 6 ++++--
 api/agent_test.go | 8 ++++----
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/api/agent.go b/api/agent.go
index 08b1441333..e2d1ad543c 100644
--- a/api/agent.go
+++ b/api/agent.go
@@ -276,9 +276,10 @@ func (a *Agent) ForceLeave(node string) error {
 
 // EnableServiceMaintenance toggles service maintenance mode on
 // for the given service ID.
-func (a *Agent) EnableServiceMaintenance(serviceID string) error {
+func (a *Agent) EnableServiceMaintenance(serviceID, reason string) error {
 	r := a.c.newRequest("PUT", "/v1/agent/service/maintenance/"+serviceID)
 	r.params.Set("enable", "true")
+	r.params.Set("reason", reason)
 	_, resp, err := requireOK(a.c.doRequest(r))
 	if err != nil {
 		return err
@@ -302,9 +303,10 @@ func (a *Agent) DisableServiceMaintenance(serviceID string) error {
 
 // EnableNodeMaintenance toggles node maintenance mode on for the
 // agent we are connected to.
-func (a *Agent) EnableNodeMaintenance() error {
+func (a *Agent) EnableNodeMaintenance(reason string) error {
 	r := a.c.newRequest("PUT", "/v1/agent/maintenance")
 	r.params.Set("enable", "true")
+	r.params.Set("reason", reason)
 	_, resp, err := requireOK(a.c.doRequest(r))
 	if err != nil {
 		return err
diff --git a/api/agent_test.go b/api/agent_test.go
index e27887d6f8..37a00c573f 100644
--- a/api/agent_test.go
+++ b/api/agent_test.go
@@ -272,7 +272,7 @@ func TestServiceMaintenance(t *testing.T) {
 	}
 
 	// Enable maintenance mode
-	if err := agent.EnableServiceMaintenance("redis"); err != nil {
+	if err := agent.EnableServiceMaintenance("redis", "broken"); err != nil {
 		t.Fatalf("err: %s", err)
 	}
 
@@ -285,7 +285,7 @@ func TestServiceMaintenance(t *testing.T) {
 	for _, check := range checks {
 		if strings.Contains(check.CheckID, "maintenance") {
 			found = true
-			if check.Status != "critical" {
+			if check.Status != "critical" || check.Notes != "broken" {
 				t.Fatalf("bad: %#v", checks)
 			}
 		}
@@ -318,7 +318,7 @@ func TestNodeMaintenance(t *testing.T) {
 	agent := c.Agent()
 
 	// Enable maintenance mode
-	if err := agent.EnableNodeMaintenance(); err != nil {
+	if err := agent.EnableNodeMaintenance("broken"); err != nil {
 		t.Fatalf("err: %s", err)
 	}
 
@@ -331,7 +331,7 @@ func TestNodeMaintenance(t *testing.T) {
 	for _, check := range checks {
 		if strings.Contains(check.CheckID, "maintenance") {
 			found = true
-			if check.Status != "critical" {
+			if check.Status != "critical" || check.Notes != "broken" {
 				t.Fatalf("bad: %#v", checks)
 			}
 		}