diff --git a/command/agent/node_endpoint.go b/command/agent/node_endpoint.go index 4d326a98c..2fdf3cfce 100644 --- a/command/agent/node_endpoint.go +++ b/command/agent/node_endpoint.go @@ -84,7 +84,7 @@ func (s *HTTPServer) nodeAllocations(resp http.ResponseWriter, req *http.Request } func (s *HTTPServer) nodeToggleDrain(resp http.ResponseWriter, req *http.Request, - jobName string) (interface{}, error) { + nodeID string) (interface{}, error) { if req.Method != "PUT" && req.Method != "POST" { return nil, CodedError(405, ErrInvalidMethod) } @@ -99,8 +99,18 @@ func (s *HTTPServer) nodeToggleDrain(resp http.ResponseWriter, req *http.Request return nil, CodedError(400, "invalid enable value") } - // TODO - return nil, nil + args := structs.NodeUpdateDrainRequest{ + NodeID: nodeID, + Drain: enable, + } + s.parseRegion(req, &args.Region) + + var out structs.NodeDrainUpdateResponse + if err := s.agent.RPC("Client.UpdateDrain", &args, &out); err != nil { + return nil, err + } + setIndex(resp, out.Index) + return out, nil } func (s *HTTPServer) nodeQuery(resp http.ResponseWriter, req *http.Request, diff --git a/command/agent/node_endpoint_test.go b/command/agent/node_endpoint_test.go index b071260a6..6240ec375 100644 --- a/command/agent/node_endpoint_test.go +++ b/command/agent/node_endpoint_test.go @@ -158,6 +158,54 @@ func TestHTTP_NodeAllocations(t *testing.T) { }) } +func TestHTTP_NodeDrain(t *testing.T) { + httpTest(t, nil, func(s *TestServer) { + // Create the node + node := mock.Node() + args := structs.NodeRegisterRequest{ + Node: node, + WriteRequest: structs.WriteRequest{Region: "region1"}, + } + var resp structs.NodeUpdateResponse + if err := s.Agent.RPC("Client.Register", &args, &resp); err != nil { + t.Fatalf("err: %v", err) + } + + // Directly manipulate the state + state := s.Agent.server.State() + alloc1 := mock.Alloc() + alloc1.NodeID = node.ID + err := state.UpdateAllocations(1000, []*structs.Allocation{alloc1}) + if err != nil { + t.Fatalf("err: %v", err) + } + + // Make the HTTP request + req, err := http.NewRequest("POST", "/v1/node/"+node.ID+"/drain?enable=1", nil) + if err != nil { + t.Fatalf("err: %v", err) + } + respW := httptest.NewRecorder() + + // Make the request + obj, err := s.Server.NodeSpecificRequest(respW, req) + if err != nil { + t.Fatalf("err: %v", err) + } + + // Check for the index + if respW.HeaderMap.Get("X-Nomad-Index") == "" { + t.Fatalf("missing index") + } + + // Check the response + upd := obj.(structs.NodeDrainUpdateResponse) + if len(upd.EvalIDs) == 0 { + t.Fatalf("bad: %v", upd) + } + }) +} + func TestHTTP_NodeQuery(t *testing.T) { httpTest(t, nil, func(s *TestServer) { // Create the job