diff --git a/command/node_drain.go b/command/node_drain.go index 229b99c64..cec157abb 100644 --- a/command/node_drain.go +++ b/command/node_drain.go @@ -3,6 +3,9 @@ package command import ( "fmt" "strings" + + "github.com/hashicorp/nomad/api/contexts" + "github.com/posener/complete" ) type NodeDrainCommand struct { @@ -42,6 +45,25 @@ func (c *NodeDrainCommand) Synopsis() string { return "Toggle drain mode on a given node" } +func (c *NodeDrainCommand) AutocompleteFlags() complete.Flags { + return nil +} + +func (c *NodeDrainCommand) AutocompleteArgs() complete.Predictor { + client, _ := c.Meta.Client() + return complete.PredictFunc(func(a complete.Args) []string { + if len(a.Completed) > 1 { + return nil + } + + resp, err := client.Search().PrefixSearch(a.Last, contexts.Nodes) + if err != nil { + return []string{} + } + return resp.Matches[contexts.Nodes] + }) +} + func (c *NodeDrainCommand) Run(args []string) int { var enable, disable, self, autoYes bool diff --git a/command/node_drain_test.go b/command/node_drain_test.go index 3bc551e17..9644641fb 100644 --- a/command/node_drain_test.go +++ b/command/node_drain_test.go @@ -1,10 +1,14 @@ package command import ( + "fmt" "strings" "testing" + "github.com/hashicorp/nomad/testutil" "github.com/mitchellh/cli" + "github.com/posener/complete" + "github.com/stretchr/testify/assert" ) func TestNodeDrainCommand_Implements(t *testing.T) { @@ -82,3 +86,45 @@ func TestNodeDrainCommand_Fails(t *testing.T) { t.Fatalf("expected not exist error, got: %s", out) } } + +func TestNodeDrainCommand_AutocompleteArgs(t *testing.T) { + assert := assert.New(t) + t.Parallel() + + srv, client, url := testServer(t, true, nil) + defer srv.Shutdown() + + // Wait for a node to appear + var nodeID string + testutil.WaitForResult(func() (bool, error) { + nodes, _, err := client.Nodes().List(nil) + if err != nil { + return false, err + } + if len(nodes) == 0 { + return false, fmt.Errorf("missing node") + } + nodeID = nodes[0].ID + return true, nil + }, func(err error) { + t.Fatalf("err: %s", err) + }) + + ui := new(cli.MockUi) + cmd := &NodeDrainCommand{Meta: Meta{Ui: ui, flagAddress: url}} + + prefix := nodeID[:len(nodeID)-5] + args := complete.Args{Last: prefix} + predictor := cmd.AutocompleteArgs() + + res := predictor.Predict(args) + assert.Equal(1, len(res)) + assert.Equal(nodeID, res[0]) + + // Autocomplete should only complete once + args = complete.Args{Last: prefix, Completed: []string{prefix, "1", "2"}} + predictor = cmd.AutocompleteArgs() + + res = predictor.Predict(args) + assert.Nil(res) +}