node drain

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
 

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)
+}