diff --git a/command/agent/agent_endpoint_test.go b/command/agent/agent_endpoint_test.go index de785e4312..d389cdf03d 100644 --- a/command/agent/agent_endpoint_test.go +++ b/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") + } +} diff --git a/command/agent/agent_test.go b/command/agent/agent_test.go index ae908e92f4..51f845f5ae 100644 --- a/command/agent/agent_test.go +++ b/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,24 +806,40 @@ func TestAgent_MaintenanceMode(t *testing.T) { } // Make sure the critical health check was added - for _, check := range agent.state.Checks() { - if check.CheckID == serviceMaintCheckID { - return - } + if _, ok := agent.state.Checks()[serviceMaintCheckID]; !ok { + t.Fatalf("should have registered critical maintenance check") } - // 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 == serviceMaintCheckID { - t.Fatalf("should have deregistered maintenance check") - } + if _, ok := agent.state.Checks()[serviceMaintCheckID]; ok { + 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") } }