diff --git a/api/agent.go b/api/agent.go index b71f2c32d..08b144133 100644 --- a/api/agent.go +++ b/api/agent.go @@ -273,3 +273,55 @@ func (a *Agent) ForceLeave(node string) error { resp.Body.Close() return nil } + +// EnableServiceMaintenance toggles service maintenance mode on +// for the given service ID. +func (a *Agent) EnableServiceMaintenance(serviceID string) error { + r := a.c.newRequest("PUT", "/v1/agent/service/maintenance/"+serviceID) + r.params.Set("enable", "true") + _, resp, err := requireOK(a.c.doRequest(r)) + if err != nil { + return err + } + resp.Body.Close() + return nil +} + +// DisableServiceMaintenance toggles service maintenance mode off +// for the given service ID. +func (a *Agent) DisableServiceMaintenance(serviceID string) error { + r := a.c.newRequest("PUT", "/v1/agent/service/maintenance/"+serviceID) + r.params.Set("enable", "false") + _, resp, err := requireOK(a.c.doRequest(r)) + if err != nil { + return err + } + resp.Body.Close() + return nil +} + +// EnableNodeMaintenance toggles node maintenance mode on for the +// agent we are connected to. +func (a *Agent) EnableNodeMaintenance() error { + r := a.c.newRequest("PUT", "/v1/agent/maintenance") + r.params.Set("enable", "true") + _, resp, err := requireOK(a.c.doRequest(r)) + if err != nil { + return err + } + resp.Body.Close() + return nil +} + +// DisableNodeMaintenance toggles node maintenance mode off for the +// agent we are connected to. +func (a *Agent) DisableNodeMaintenance() error { + r := a.c.newRequest("PUT", "/v1/agent/maintenance") + r.params.Set("enable", "false") + _, resp, err := requireOK(a.c.doRequest(r)) + if err != nil { + return err + } + resp.Body.Close() + return nil +} diff --git a/api/agent_test.go b/api/agent_test.go index 627256c23..a6e2fc542 100644 --- a/api/agent_test.go +++ b/api/agent_test.go @@ -255,3 +255,57 @@ func TestAgent_ForceLeave(t *testing.T) { t.Fatalf("err: %v", err) } } + +func TestServiceMaintenance(t *testing.T) { + c, s := makeClient(t) + defer s.stop() + + agent := c.Agent() + + // First register a service + serviceReg := &AgentServiceRegistration{ + Name: "redis", + } + if err := agent.ServiceRegister(serviceReg); err != nil { + t.Fatalf("err: %v", err) + } + + // Enable maintenance mode + if err := agent.EnableServiceMaintenance("redis"); err != nil { + t.Fatalf("err: %s", err) + } + + // Ensure a critical check was added + checks, err := agent.Checks() + if err != nil { + t.Fatalf("err: %v", err) + } + found := false + for _, check := range checks { + if check.ServiceName == "redis" { + found = true + if check.Status != "critical" { + t.Fatalf("bad: %#v", checks) + } + } + } + if !found { + t.Fatalf("bad: %#v", checks) + } + + // Disable maintenance mode + if err := agent.DisableServiceMaintenance("redis"); err != nil { + t.Fatalf("err: %s", err) + } + + // Ensure the critical health check was removed + checks, err = agent.Checks() + if err != nil { + t.Fatalf("err: %s", err) + } + for _, check := range checks { + if check.ServiceID == "redis" { + t.Fatalf("should have removed health check") + } + } +}