commandkit/internal/cobra/completion_test.go

package cobra

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"testing"

	spfcobra "github.com/spf13/cobra"
)

func TestCompletionScriptGeneration(t *testing.T) {
	for _, shell := range []string{"bash", "zsh", "fish"} {
		var output bytes.Buffer
		if err := CompletionScript(&output, "sample-app", shell); err != nil {
			t.Fatalf("%s: %v", shell, err)
		}
		if !strings.Contains(output.String(), "__complete") {
			t.Errorf("%s script lacks a completion request", shell)
		}
	}
	if err := CompletionScript(io.Discard, "app", "other"); err == nil {
		t.Fatal("unsupported shell succeeded")
	}
	if err := CompletionScript(shortCompletionWriter{}, "app", "bash"); !errors.Is(err, io.ErrShortWrite) {
		t.Fatalf("short write: %v", err)
	}
}

type shortCompletionWriter struct{}

func (shortCompletionWriter) Write(p []byte) (int, error) { return len(p) - 1, nil }

// TestCompletionShells is opt-in because it requires three shells and pexpect.
func TestCompletionShells(t *testing.T) {
	if os.Getenv("COMMANDKIT_COMPLETION_SHELLS") != "1" {
		t.Skip("set COMMANDKIT_COMPLETION_SHELLS=1 to run shell completion checks")
	}
	dir := t.TempDir()
	fixture := `#!/usr/bin/python3
import json, os, sys
with open(os.environ['PROBE_LOG'], 'a') as log:
    log.write(json.dumps(sys.argv[1:]) + '\n')
sys.stdout.write(os.environ['PROBE_REPLY'])
`
	if err := os.WriteFile(filepath.Join(dir, "probe"), []byte(fixture), 0o755); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(dir, "probe.app"), []byte(fixture), 0o755); err != nil {
		t.Fatal(err)
	}
	if err := os.Mkdir(filepath.Join(dir, "partial dir"), 0o755); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(dir, "partial file"), nil, 0o644); err != nil {
		t.Fatal(err)
	}
	reply := func(candidates string, directive spfcobra.ShellCompDirective) string {
		return fmt.Sprintf("%s:%d\n", candidates, directive)
	}
	tests := []struct {
		name, line, reply string
		query, executed   []string
	}{
		{"attached choice", "probe --config=./par", reply("./partial file\tdescription\n", spfcobra.ShellCompDirectiveNoFileComp), []string{"__complete", "--config=./par"}, []string{"--config=./partial file"}},
		{"choice with equals", "probe --config=foo=ba", reply("foo=bar\tdescription\n", spfcobra.ShellCompDirectiveNoFileComp), []string{"__complete", "--config=foo=ba"}, []string{"--config=foo=bar"}},
		{"separated file", `probe --config ./partial\ f`, reply("", spfcobra.ShellCompDirectiveDefault), []string{"__complete", "--config", "./partial f"}, []string{"--config", "./partial file"}},
		{"attached file", `probe --config=./partial\ f`, reply("", spfcobra.ShellCompDirectiveDefault), []string{"__complete", "--config=./partial f"}, []string{"--config=./partial file"}},
		{"directory", "probe ./partial", reply("", spfcobra.ShellCompDirectiveFilterDirs), []string{"__complete", "./partial"}, []string{"./partial dir"}},
		{"attached directory", "probe --config=./partial", reply("", spfcobra.ShellCompDirectiveFilterDirs), []string{"__complete", "--config=./partial"}, []string{"--config=./partial dir"}},
		{"empty", "probe ", reply("", spfcobra.ShellCompDirectiveNoFileComp), []string{"__complete", ""}, []string{}},
		{"no file fallback", `probe ./partial\ f`, reply("", spfcobra.ShellCompDirectiveNoFileComp), []string{"__complete", "./partial f"}, []string{"./partial f"}},
		{"cluster", "probe -ab", reply("-abc\tflag\n", spfcobra.ShellCompDirectiveNoFileComp), []string{"__complete", "-ab"}, []string{"-abc"}},
		{"quoted argument", "probe 'alpha beta' --config=./par", reply("./partial file\tdescription\n", spfcobra.ShellCompDirectiveNoFileComp), []string{"__complete", "alpha beta", "--config=./par"}, []string{"alpha beta", "--config=./partial file"}},
	}
	for _, shell := range []string{"bash", "zsh", "fish"} {
		script := filepath.Join(dir, shell+".completion")
		var generated bytes.Buffer
		if err := CompletionScript(&generated, "probe", shell); err != nil {
			t.Fatal(err)
		}
		if err := os.WriteFile(script, generated.Bytes(), 0o644); err != nil {
			t.Fatal(err)
		}
		for _, tc := range tests {
			t.Run(shell+"/"+tc.name, func(t *testing.T) {
				log := filepath.Join(dir, "log")
				if err := os.WriteFile(log, nil, 0o600); err != nil {
					t.Fatal(err)
				}
				cmd := exec.Command("python3", "testdata/completion_terminal.py", shell, script, tc.line)
				cmd.Dir = "."
				cmd.Env = append(os.Environ(), "PATH="+dir+string(os.PathListSeparator)+os.Getenv("PATH"), "HOME="+dir, "TERM=xterm", "PROBE_LOG="+log, "PROBE_REPLY="+tc.reply, "PROBE_DIR="+dir)
				if out, err := cmd.CombinedOutput(); err != nil {
					t.Fatalf("shell completion: %v\n%s", err, out)
				}
				contents, err := os.ReadFile(log)
				if err != nil {
					t.Fatal(err)
				}
				var entries [][]string
				for _, line := range strings.Split(strings.TrimSpace(string(contents)), "\n") {
					var args []string
					if err := json.Unmarshal([]byte(line), &args); err != nil {
						t.Fatalf("invalid fixture log %q: %v", contents, err)
					}
					entries = append(entries, args)
				}
				if len(entries) < 2 || !equalCompletionArgs(entries[len(entries)-2], tc.query, false) || !equalCompletionArgs(entries[len(entries)-1], tc.executed, strings.Contains(tc.name, "directory")) {
					t.Fatalf("query and edited line: %q", contents)
				}
			})
		}
		t.Run(shell+"/dotted executable", func(t *testing.T) {
			var generated bytes.Buffer
			if err := CompletionScript(&generated, "probe.app", shell); err != nil {
				t.Fatal(err)
			}
			if err := os.WriteFile(script, generated.Bytes(), 0o644); err != nil {
				t.Fatal(err)
			}
			log := filepath.Join(dir, "log")
			if err := os.WriteFile(log, nil, 0o600); err != nil {
				t.Fatal(err)
			}
			cmd := exec.Command("python3", "testdata/completion_terminal.py", shell, script, "probe.app --config=./par")
			cmd.Env = append(os.Environ(), "PATH="+dir+string(os.PathListSeparator)+os.Getenv("PATH"), "HOME="+dir, "TERM=xterm", "PROBE_LOG="+log, "PROBE_REPLY="+reply("./partial file\tdescription\n", spfcobra.ShellCompDirectiveNoFileComp), "PROBE_DIR="+dir)
			if out, err := cmd.CombinedOutput(); err != nil {
				t.Fatalf("shell completion: %v\n%s", err, out)
			}
			contents, err := os.ReadFile(log)
			if err != nil || !strings.Contains(string(contents), `["__complete", "--config=./par"]`) || !strings.Contains(string(contents), `["--config=./partial file"]`) {
				t.Fatalf("query and edited line: %q; read error: %v", contents, err)
			}
		})
	}
}

func equalCompletionArgs(got, want []string, directory bool) bool {
	if len(got) != len(want) {
		return false
	}
	for i := range got {
		value := got[i]
		if directory && i == len(got)-1 {
			value = strings.TrimSuffix(value, "/")
		}
		if value != want[i] {
			return false
		}
	}
	return true
}