agent: test node maintenance mode

command/agent/agent_endpoint_test.go

@@ -607,3 +607,77 @@ func TestHTTPAgent_DisableServiceMaintenance(t *testing.T) {
 		t.Fatalf("should have removed maintenance check")
 	}
 }
+
+func TestHTTPAgent_NodeMaintenanceEndpoint_BadRequest(t *testing.T) {
+	dir, srv := makeHTTPServer(t)
+	defer os.RemoveAll(dir)
+	defer srv.Shutdown()
+	defer srv.agent.Shutdown()
+
+	// Fails on non-PUT
+	req, _ := http.NewRequest("GET", "/v1/agent/self/maintenance?enable=true", nil)
+	resp := httptest.NewRecorder()
+	if _, err := srv.AgentNodeMaintenance(resp, req); err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	if resp.Code != 405 {
+		t.Fatalf("expected 405, got %d", resp.Code)
+	}
+
+	// Fails when no enable flag provided
+	req, _ = http.NewRequest("PUT", "/v1/agent/self/maintenance", nil)
+	resp = httptest.NewRecorder()
+	if _, err := srv.AgentNodeMaintenance(resp, req); err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	if resp.Code != 400 {
+		t.Fatalf("expected 400, got %d", resp.Code)
+	}
+}
+
+func TestHTTPAgent_EnableNodeMaintenance(t *testing.T) {
+	dir, srv := makeHTTPServer(t)
+	defer os.RemoveAll(dir)
+	defer srv.Shutdown()
+	defer srv.agent.Shutdown()
+
+	// Force the node into maintenance mode
+	req, _ := http.NewRequest("PUT", "/v1/agent/self/maintenance?enable=true", nil)
+	resp := httptest.NewRecorder()
+	if _, err := srv.AgentNodeMaintenance(resp, req); err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	if resp.Code != 200 {
+		t.Fatalf("expected 200, got %d", resp.Code)
+	}
+
+	// Ensure the maintenance check was registered
+	if _, ok := srv.agent.state.Checks()[nodeMaintCheckID]; !ok {
+		t.Fatalf("should have registered maintenance check")
+	}
+}
+
+func TestHTTPAgent_DisableNodeMaintenance(t *testing.T) {
+	dir, srv := makeHTTPServer(t)
+	defer os.RemoveAll(dir)
+	defer srv.Shutdown()
+	defer srv.agent.Shutdown()
+
+	// Force the node into maintenance mode
+	srv.agent.EnableNodeMaintenance()
+
+	// Leave maintenance mode
+	req, _ := http.NewRequest("PUT", "/v1/agent/self/maintenance?enable=false", nil)
+	resp := httptest.NewRecorder()
+	if _, err := srv.AgentNodeMaintenance(resp, req); err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	if resp.Code != 200 {
+		t.Fatalf("expected 200, got %d", resp.Code)
+	}
+
+	// Ensure the maintenance check was removed
+	if _, ok := srv.agent.state.Checks()[nodeMaintCheckID]; ok {
+		t.Fatalf("should have removed maintenance check")
+	}
+}

command/agent/agent_test.go

@@ -782,7 +782,7 @@ func TestAgent_unloadServices(t *testing.T) {
 	}
 }
 
-func TestAgent_MaintenanceMode(t *testing.T) {
+func TestAgent_ServiceMaintenanceMode(t *testing.T) {
 	config := nextConfig()
 	dir, agent := makeAgent(t, config)
 	defer os.RemoveAll(dir)
@@ -806,14 +806,9 @@ func TestAgent_MaintenanceMode(t *testing.T) {
 	}
 
 	// Make sure the critical health check was added
-	for _, check := range agent.state.Checks() {
+	if _, ok := agent.state.Checks()[serviceMaintCheckID]; !ok {
-		if check.CheckID == serviceMaintCheckID {
-			return
-		}
-	}
-
-	// Didn't find the check
 		t.Fatalf("should have registered critical maintenance check")
+	}
 
 	// Leave maintenance mode
 	if err := agent.DisableServiceMaintenance("redis"); err != nil {
@@ -821,9 +816,30 @@ func TestAgent_MaintenanceMode(t *testing.T) {
 	}
 
 	// Ensure the check was deregistered
-	for _, check := range agent.state.Checks() {
+	if _, ok := agent.state.Checks()[serviceMaintCheckID]; ok {
-		if check.CheckID == serviceMaintCheckID {
 		t.Fatalf("should have deregistered maintenance check")
 	}
+}
+
+func TestAgent_NodeMaintenanceMode(t *testing.T) {
+	config := nextConfig()
+	dir, agent := makeAgent(t, config)
+	defer os.RemoveAll(dir)
+	defer agent.Shutdown()
+
+	// Enter maintenance mode for the node
+	agent.EnableNodeMaintenance()
+
+	// Make sure the critical health check was added
+	if _, ok := agent.state.Checks()[nodeMaintCheckID]; !ok {
+		t.Fatalf("should have registered critical node check")
+	}
+
+	// Leave maintenance mode
+	agent.DisableNodeMaintenance()
+
+	// Ensure the check was deregistered
+	if _, ok := agent.state.Checks()[nodeMaintCheckID]; ok {
+		t.Fatalf("should have deregistered critical node check")
 	}
 }