From 7d4801db296cb1c24d79aafb7cc6e231c3623741 Mon Sep 17 00:00:00 2001
From: Ryan Uber <ru@ryanuber.com>
Date: Thu, 15 Jan 2015 00:25:36 -0800
Subject: [PATCH] agent: test agent service maintenance mode

---
 command/agent/agent.go      |  4 ++++
 command/agent/agent_test.go | 46 +++++++++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+)

diff --git a/command/agent/agent.go b/command/agent/agent.go
index dfdf10cf3..dd374123e 100644
--- a/command/agent/agent.go
+++ b/command/agent/agent.go
@@ -999,6 +999,8 @@ func (a *Agent) unloadChecks() error {
 	return nil
 }
 
+// EnableServiceMaintenance will register a false health check against the given
+// service ID with critical status. This will exclude the service from queries.
 func (a *Agent) EnableServiceMaintenance(serviceID string) error {
 	var service *structs.NodeService
 	for _, svc := range a.state.Services() {
@@ -1034,6 +1036,8 @@ func (a *Agent) EnableServiceMaintenance(serviceID string) error {
 	return nil
 }
 
+// DisableServiceMaintenance will deregister the fake maintenance mode check
+// if the service has been marked as in maintenance.
 func (a *Agent) DisableServiceMaintenance(serviceID string) error {
 	var service *structs.NodeService
 	for _, svc := range a.state.Services() {
diff --git a/command/agent/agent_test.go b/command/agent/agent_test.go
index 0add36702..7b2ec3b71 100644
--- a/command/agent/agent_test.go
+++ b/command/agent/agent_test.go
@@ -781,3 +781,49 @@ func TestAgent_unloadServices(t *testing.T) {
 		t.Fatalf("consul service should not be removed")
 	}
 }
+
+func TestAgent_MaintenanceMode(t *testing.T) {
+	config := nextConfig()
+	dir, agent := makeAgent(t, config)
+	defer os.RemoveAll(dir)
+	defer agent.Shutdown()
+
+	svc := &structs.NodeService{
+		ID:      "redis",
+		Service: "redis",
+		Tags:    []string{"foo"},
+		Port:    8000,
+	}
+
+	// Register the service
+	if err := agent.AddService(svc, nil, false); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	// Enter maintenance mode for the service
+	if err := agent.EnableServiceMaintenance("redis"); err != nil {
+		t.Fatalf("err: %s", err)
+	}
+
+	// Make sure the critical health check was added
+	for _, check := range agent.state.Checks() {
+		if check.CheckID == maintCheckID {
+			return
+		}
+	}
+
+	// Didn't find the check
+	t.Fatalf("should have registered critical maintenance check")
+
+	// Leave maintenance mode
+	if err := agent.DisableServiceMaintenance("redis"); err != nil {
+		t.Fatalf("err: %s", err)
+	}
+
+	// Ensure the check was deregistered
+	for _, check := range agent.state.Checks() {
+		if check.CheckID == maintCheckID {
+			t.Fatalf("should have deregistered maintenance check")
+		}
+	}
+}