add some tests

command/alloc_exec.go

@@ -19,6 +19,10 @@ import (
 
 type AllocExecCommand struct {
 	Meta
+
+	Stdin  io.Reader
+	Stdout io.WriteCloser
+	Stderr io.WriteCloser
 }
 
 func (l *AllocExecCommand) Help() string {
@@ -126,7 +130,7 @@ func (l *AllocExecCommand) Run(args []string) int {
 		l.Ui.Error(commandErrorText(l))
 		return 1
 	} else if numArgs < 2 {
-		l.Ui.Error("This command takes command as arguments")
+		l.Ui.Error("A command is required")
 		l.Ui.Error(commandErrorText(l))
 		return 1
 	}
@@ -205,12 +209,22 @@ func (l *AllocExecCommand) Run(args []string) int {
 		}
 	}
 
-	var stdin io.Reader = os.Stdin
+	if l.Stdin == nil {
+		l.Stdin = os.Stdin
+	}
+	if l.Stdout == nil {
+		l.Stdout = os.Stdout
+	}
+	if l.Stderr == nil {
+		l.Stderr = os.Stderr
+	}
+
+	var stdin io.Reader = l.Stdin
 	if !stdinOpt {
 		stdin = bytes.NewReader(nil)
 	}
 
-	code, err := l.execImpl(client, alloc, task, ttyOpt, command, stdin, os.Stdout, os.Stderr)
+	code, err := l.execImpl(client, alloc, task, ttyOpt, command, stdin, l.Stdout, l.Stderr)
 	if err != nil {
 		l.Ui.Error(fmt.Sprintf("failed to exec into task: %v", err))
 		return 1

command/alloc_exec_test.go

+package command
+
+import (
+	"bytes"
+	"fmt"
+	"testing"
+
+	"github.com/hashicorp/nomad/helper/uuid"
+	"github.com/hashicorp/nomad/nomad/mock"
+	"github.com/hashicorp/nomad/nomad/structs"
+	"github.com/hashicorp/nomad/testutil"
+	"github.com/mitchellh/cli"
+	"github.com/posener/complete"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// static check
+var _ cli.Command = &AllocExecCommand{}
+
+func TestAllocExecCommand_Fails(t *testing.T) {
+	t.Parallel()
+	srv, _, url := testServer(t, false, nil)
+	defer srv.Shutdown()
+
+	cases := []struct {
+		name          string
+		args          []string
+		expectedError string
+	}{
+		{
+			"misuse",
+			[]string{"bad"},
+			commandErrorText(&AllocExecCommand{}),
+		},
+		{
+			"connection failure",
+			[]string{"-address=nope", "26470238-5CF2-438F-8772-DC67CFB0705C", "/bin/bash"},
+			"Error querying allocation",
+		},
+		{
+			"not found alloc",
+			[]string{"-address=" + url, "26470238-5CF2-438F-8772-DC67CFB0705C", "/bin/bash"},
+			"No allocation(s) with prefix or id",
+		},
+		{
+			"too short allocis",
+			[]string{"-address=" + url, "2", "/bin/bash"},
+			"Alloc ID must contain at least two characters",
+		},
+		{
+			"missing command",
+			[]string{"-address=" + url, "26470238-5CF2-438F-8772-DC67CFB0705C"},
+			"A command is required",
+		},
+	}
+
+	for _, c := range cases {
+		t.Run(c.name, func(t *testing.T) {
+			ui := new(cli.MockUi)
+			cmd := &AllocExecCommand{Meta: Meta{Ui: ui}}
+
+			code := cmd.Run(c.args)
+			require.Equal(t, 1, code)
+
+			require.Contains(t, ui.ErrorWriter.String(), c.expectedError)
+
+			ui.ErrorWriter.Reset()
+			ui.OutputWriter.Reset()
+
+		})
+	}
+}
+
+func TestAllocExecCommand_AutocompleteArgs(t *testing.T) {
+	assert := assert.New(t)
+	t.Parallel()
+
+	srv, _, url := testServer(t, true, nil)
+	defer srv.Shutdown()
+
+	ui := new(cli.MockUi)
+	cmd := &AllocExecCommand{Meta: Meta{Ui: ui, flagAddress: url}}
+
+	// Create a fake alloc
+	state := srv.Agent.Server().State()
+	a := mock.Alloc()
+	assert.Nil(state.UpsertAllocs(1000, []*structs.Allocation{a}))
+
+	prefix := a.ID[:5]
+	args := complete.Args{Last: prefix}
+	predictor := cmd.AutocompleteArgs()
+
+	res := predictor.Predict(args)
+	assert.Equal(1, len(res))
+	assert.Equal(a.ID, res[0])
+}
+
+func TestAllocExecCommand_Run(t *testing.T) {
+	t.Parallel()
+	srv, client, url := testServer(t, true, nil)
+	defer srv.Shutdown()
+
+	// Wait for a node to be ready
+	testutil.WaitForResult(func() (bool, error) {
+		nodes, _, err := client.Nodes().List(nil)
+		if err != nil {
+			return false, err
+		}
+
+		for _, node := range nodes {
+			if _, ok := node.Drivers["mock_driver"]; ok &&
+				node.Status == structs.NodeStatusReady {
+				return true, nil
+			}
+		}
+		return false, fmt.Errorf("no ready nodes")
+	}, func(err error) {
+		require.NoError(t, err)
+	})
+
+	jobID := uuid.Generate()
+	job := testJob(jobID)
+	job.TaskGroups[0].Tasks[0].Config = map[string]interface{}{
+		"run_for": "10s",
+		"exec_command": map[string]interface{}{
+			"run_for":       "1ms",
+			"exit_code":     21,
+			"stdout_string": "sample stdout output\n",
+			"stderr_string": "sample stderr output\n",
+		},
+	}
+	resp, _, err := client.Jobs().Register(job, nil)
+	require.NoError(t, err)
+
+	evalUi := new(cli.MockUi)
+	code := waitForSuccess(evalUi, client, fullId, t, resp.EvalID)
+	require.Equal(t, 0, code, "failed to get status - output: %v", evalUi.ErrorWriter.String())
+
+	// get an alloc id
+	allocId := ""
+	if allocs, _, err := client.Jobs().Allocations(jobID, false, nil); err == nil {
+		if len(allocs) > 0 {
+			allocId = allocs[0].ID
+		}
+	}
+	require.NotEmpty(t, allocId, "failed to find allocation")
+
+	ui := new(cli.MockUi)
+	var stdout, stderr bufferCloser
+
+	cmd := &AllocExecCommand{
+		Meta:   Meta{Ui: ui},
+		Stdin:  bytes.NewReader(nil),
+		Stdout: &stdout,
+		Stderr: &stderr,
+	}
+
+	code = cmd.Run([]string{"-address=" + url, allocId, "simpelcommand"})
+	assert.Equal(t, 21, code)
+	assert.Contains(t, stdout.String(), "sample stdout output")
+	assert.Contains(t, stderr.String(), "sample stderr output")
+}
+
+type bufferCloser struct {
+	bytes.Buffer
+}
+
+func (b *bufferCloser) Close() error {
+	return nil
+}