Fix shell completion

pkg/app/commands.go

@@ -46,6 +46,7 @@ import (
 	"github.com/fastly/cli/pkg/commands/purge"
 	"github.com/fastly/cli/pkg/commands/service"
 	"github.com/fastly/cli/pkg/commands/serviceversion"
+	"github.com/fastly/cli/pkg/commands/shellcomplete"
 	"github.com/fastly/cli/pkg/commands/stats"
 	"github.com/fastly/cli/pkg/commands/update"
 	"github.com/fastly/cli/pkg/commands/user"
@@ -63,6 +64,7 @@ func defineCommands(
 	globals *config.Data,
 	data manifest.Data,
 	opts RunOpts) []cmd.Command {
+	shellcompleteCmdRoot := shellcomplete.NewRootCommand(app, globals)
 	aclCmdRoot := acl.NewRootCommand(app, globals)
 	aclCreate := acl.NewCreateCommand(aclCmdRoot.CmdClause, globals, data)
 	aclDelete := acl.NewDeleteCommand(aclCmdRoot.CmdClause, globals, data)
@@ -326,6 +328,7 @@ func defineCommands(
 	whoamiCmdRoot := whoami.NewRootCommand(app, opts.HTTPClient, globals)
 
 	return []cmd.Command{
+		shellcompleteCmdRoot,
 		aclCmdRoot,
 		aclCreate,
 		aclDelete,

pkg/app/run.go

@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"io"
 	"os"
-	"regexp"
 	"time"
 
 	"github.com/fastly/cli/pkg/api"
@@ -21,10 +20,6 @@ import (
 	"github.com/fastly/kingpin"
 )
 
-var (
-	completionRegExp = regexp.MustCompile("completion-(?:script-)?(?:bash|zsh)$")
-)
-
 // Versioners represents all supported versioner types.
 type Versioners struct {
 	CLI     update.Versioner
@@ -81,12 +76,6 @@ func Run(opts RunOpts) error {
 	// error states and output control flow.
 	app.Terminate(nil)
 
-	// As kingpin generates bash completion as a side-effect of kingpin.Parse we
-	// allow it to call os.Exit, only if a completetion flag is present.
-	if isCompletion(opts.Args) {
-		app.Terminate(os.Exit)
-	}
-
 	// WARNING: kingping has no way of decorating flags as being "global"
 	// therefore if you add/remove a global flag you will also need to update
 	// the globalFlag map in pkg/app/usage.go which is used for usage rendering.
@@ -105,8 +94,16 @@ func Run(opts RunOpts) error {
 	if err != nil {
 		return err
 	}
-	// We add a special case for when cmd.ArgsIsHelpJSON() is true.
-	if name == "help--format=json" || name == "help--formatjson" {
+	// We short-circuit the execution for specific cases:
+	//
+	// - cmd.ArgsIsHelpJSON() == true
+	// - shell autocompletion flag provided
+	switch name {
+	case "help--format=json":
+		fallthrough
+	case "help--formatjson":
+		fallthrough
+	case "shell-autocomplete":
 		return nil
 	}
 
@@ -185,15 +182,3 @@ func FastlyAPIClient(token, endpoint string) (api.Interface, error) {
 	client, err := fastly.NewClientForEndpoint(token, endpoint)
 	return client, err
 }
-
-// isCompletion determines whether the supplied command arguments are for
-// bash/zsh completion output.
-func isCompletion(args []string) bool {
-	var found bool
-	for _, arg := range args {
-		if completionRegExp.MatchString(arg) {
-			found = true
-		}
-	}
-	return found
-}

pkg/app/run_test.go

@@ -3,6 +3,8 @@ package app_test
 import (
 	"bufio"
 	"bytes"
+	"io"
+	"os"
 	"strings"
 	"testing"
 
@@ -45,6 +47,7 @@ func TestApplication(t *testing.T) {
 				stdout bytes.Buffer
 				stderr bytes.Buffer
 			)
+
 			opts := testutil.NewRunOpts(testcase.Args, &stdout)
 			err := app.Run(opts)
 			if err != nil {
@@ -56,6 +59,84 @@ func TestApplication(t *testing.T) {
 	}
 }
 
+func TestShellCompletion(t *testing.T) {
+	args := testutil.Args
+	scenarios := []testutil.TestScenario{
+		{
+			Name: "bash shell complete",
+			Args: args("--completion-script-bash"),
+			WantOutput: `
+_fastly_bash_autocomplete() {
+    local cur prev opts base
+    COMPREPLY=()
+    cur="${COMP_WORDS[COMP_CWORD]}"
+    opts=$( ${COMP_WORDS[0]} --completion-bash ${COMP_WORDS[@]:1:$COMP_CWORD} )
+    COMPREPLY=( $(compgen -W "${opts}" -- ${cur}) )
+    return 0
+}
+complete -F _fastly_bash_autocomplete fastly
+
+`,
+		},
+		{
+			Name: "zsh shell complete",
+			Args: args("--completion-script-zsh"),
+			WantOutput: `
+#compdef fastly
+autoload -U compinit && compinit
+autoload -U bashcompinit && bashcompinit
+
+_fastly_bash_autocomplete() {
+    local cur prev opts base
+    COMPREPLY=()
+    cur="${COMP_WORDS[COMP_CWORD]}"
+    opts=$( ${COMP_WORDS[0]} --completion-bash ${COMP_WORDS[@]:1:$COMP_CWORD} )
+    COMPREPLY=( $(compgen -W "${opts}" -- ${cur}) )
+    [[ $COMPREPLY ]] && return
+    compgen -f
+    return 0
+}
+complete -F _fastly_bash_autocomplete fastly
+`,
+		},
+	}
+	for _, testcase := range scenarios {
+		t.Run(testcase.Name, func(t *testing.T) {
+			var (
+				stdout bytes.Buffer
+				stderr bytes.Buffer
+			)
+
+			// NOTE: The Kingpin dependency internally overrides our stdout
+			// variable when doing shell completion to the os.Stdout variable and so
+			// in order for us to verify it contains the shell completion output, we
+			// need an os.Pipe so we can copy off anything written to os.Stdout.
+			old := os.Stdout
+			r, w, _ := os.Pipe()
+			os.Stdout = w
+			outC := make(chan string)
+
+			go func() {
+				var buf bytes.Buffer
+				io.Copy(&buf, r)
+				outC <- buf.String()
+			}()
+
+			opts := testutil.NewRunOpts(testcase.Args, &stdout)
+			err := app.Run(opts)
+			if err != nil {
+				errors.Deduce(err).Print(&stderr)
+			}
+
+			w.Close()
+			os.Stdout = old
+			out := <-outC
+
+			testutil.AssertString(t, testcase.WantOutput, stripTrailingSpace(out))
+		})
+	}
+}
+
 // stripTrailingSpace removes any trailing spaces from the multiline str.
 func stripTrailingSpace(str string) string {
 	buf := bytes.NewBuffer(nil)

pkg/app/usage.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"os"
 	"strings"
 	"text/template"
 
@@ -353,8 +354,12 @@ func processCommandInput(
 	// be set to `help [<command>...]` as it's a built-in command that we don't
 	// control, and the latter --help flag variation will be an empty string as
 	// there were no actual 'command' specified.
+	//
+	// Additionally we don't want to use the ctx if dealing with a shell
+	// completion flag, as that depends on kingpin.Parse() being called, and so
+	// the ctx is otherwise empty.
 	var found bool
-	if !cmd.IsHelpOnly(opts.Args) && !cmd.IsHelpFlagOnly(opts.Args) {
+	if !cmd.IsHelpOnly(opts.Args) && !cmd.IsHelpFlagOnly(opts.Args) && !cmd.IsCompletion(opts.Args) {
 		command, found = cmd.Select(ctx.SelectedCommand.FullCommand(), commands)
 		if !found {
 			return command, cmdName, help(vars, err)
@@ -365,6 +370,25 @@ func processCommandInput(
 		return command, cmdName, help(vars, nil)
 	}
 
+	// NOTE: app.Parse() resets the default values for app.Writers() from
+	// io.Discard to os.Stdout and os.Stderr, meaning when using a shell
+	// autocomplete flag we'll not only see the expected output but also a help
+	// message because the parser has no matching command and so it thinks there
+	// is an error and prints the help output for us.
+	//
+	// The only way I've found to prevent this is by ensuring the arguments
+	// provided have a valid command along with the flag, for example:
+	//
+	// fastly --completion-script-bash acl
+	//
+	// But rather than rely on a feature command, we have defined a hidden
+	// command that we can safely append to the arguments and not have to worry
+	// about it getting removed accidentally in the future as we now have a test
+	// to validate the shell autocomplete behaviours.
+	if cmd.IsCompletion(opts.Args) {
+		opts.Args = append(opts.Args, "shellcomplete")
+	}
+
 	cmdName, err = app.Parse(opts.Args)
 	if err != nil {
 		return command, cmdName, help(vars, err)
@@ -373,6 +397,13 @@ func processCommandInput(
 	// Restore output writers
 	app.Writers(opts.Stdout, io.Discard)
 
+	// Kingpin generates shell completion as a side-effect of kingpin.Parse() so
+	// we allow it to call os.Exit, only if a completion flag is present.
+	if cmd.IsCompletion(opts.Args) {
+		app.Terminate(os.Exit)
+		return command, "shell-autocomplete", nil
+	}
+
 	// A side-effect of suppressing app.Parse from writing output is the usage
 	// isn't printed for the default `help` command. Therefore we capture it
 	// here by calling Parse, again swapping the Writers. This also ensures the

pkg/cmd/flags.go

@@ -5,6 +5,7 @@ import (
 	"io"
 	"os"
 	"path/filepath"
+	"regexp"
 	"sort"
 	"strconv"
 	"strings"
@@ -16,6 +17,10 @@ import (
 	"github.com/fastly/kingpin"
 )
 
+var (
+	completionRegExp = regexp.MustCompile("completion-(?:script-)?(?:bash|zsh)$")
+)
+
 // RegisterServiceIDFlag defines a --service-id flag that will attempt to
 // acquire the Service ID from multiple sources.
 //
@@ -189,3 +194,15 @@ func ContextHasHelpFlag(ctx *kingpin.ParseContext) bool {
 	_, ok := ctx.Elements.FlagMap()["help"]
 	return ok
 }
+
+// IsCompletion determines whether the supplied command arguments are for
+// bash/zsh completion output.
+func IsCompletion(args []string) bool {
+	var found bool
+	for _, arg := range args {
+		if completionRegExp.MatchString(arg) {
+			found = true
+		}
+	}
+	return found
+}

pkg/commands/shellcomplete/root.go

@@ -0,0 +1,28 @@
+package shellcomplete
+
+import (
+	"io"
+
+	"github.com/fastly/cli/pkg/cmd"
+	"github.com/fastly/cli/pkg/config"
+)
+
+// RootCommand is the parent command for all subcommands in this package.
+// It should be installed under the primary root command.
+type RootCommand struct {
+	cmd.Base
+	// no flags
+}
+
+// NewRootCommand returns a new command registered in the parent.
+func NewRootCommand(parent cmd.Registerer, globals *config.Data) *RootCommand {
+	var c RootCommand
+	c.Globals = globals
+	c.CmdClause = parent.Command("shellcomplete", "Hidden command used to prevent help output when using --completion-script-<T>").Hidden()
+	return &c
+}
+
+// Exec implements the command interface.
+func (c *RootCommand) Exec(in io.Reader, out io.Writer) error {
+	panic("unreachable")
+}