commandkit/internal/cobra/completion_test.go

package cobra

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

	commandmeta "gitea.maximumdirect.net/eric/CommandKit/internal/commandmeta"
	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)
	}
}

func TestCompletionReplyFiltersUnrepresentableValues(t *testing.T) {
	result := commandmeta.CompletionResult{Kind: commandmeta.CompletionChoices, Prefix: "-c", Choices: []commandmeta.CompletionChoice{
		{Value: "good", Description: "one\ntwo\tthree"}, {Value: ""}, {Value: "bad\nline"}, {Value: "bad\tfield"}, {Value: "bad\x00byte"},
	}}
	reply, err := CompletionReply(result)
	want := fmt.Sprintf("%s-c\ngood\tone two three\n:%d\n", completionPathPrefixMarker, spfcobra.ShellCompDirectiveNoFileComp)
	if err != nil || reply != want {
		t.Fatalf("reply=%q error=%v", reply, err)
	}
	if _, err := CompletionReply(commandmeta.CompletionResult{Kind: 255}); err == nil {
		t.Fatal("unsupported strategy accepted")
	}
}

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)
	}
	for _, executable := range []string{"probe.app", "café"} {
		if err := os.WriteFile(filepath.Join(dir, executable), []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)
	}
	if err := os.WriteFile(filepath.Join(dir, "foo=bar"), nil, 0o644); err != nil {
		t.Fatal(err)
	}
	reply := func(candidates string, directive spfcobra.ShellCompDirective) string {
		return fmt.Sprintf("%s:%d\n", candidates, directive)
	}
	pathReply := func(kind commandmeta.CompletionKind, prefix string) string {
		result, err := CompletionReply(commandmeta.CompletionResult{Kind: kind, Prefix: prefix})
		if err != nil {
			t.Fatal(err)
		}
		return result
	}
	choiceReply := func(prefix, value string) string {
		result, err := CompletionReply(commandmeta.CompletionResult{Kind: commandmeta.CompletionChoices, Prefix: prefix, Choices: []commandmeta.CompletionChoice{{Value: value, Description: "value"}}})
		if err != nil {
			t.Fatal(err)
		}
		return result
	}
	tests := []struct {
		name, line, reply string
		query, executed   []string
	}{
		{"attached choice", "probe --config=./par", choiceReply("--config=", "./partial file"), []string{"__complete", "--config=./par"}, []string{"--config=./partial file"}},
		{"choice with equals", "probe --config=foo=ba", choiceReply("--config=", "foo=bar"), []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`, pathReply(commandmeta.CompletionFiles, "--config="), []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", pathReply(commandmeta.CompletionDirectories, "--config="), []string{"__complete", "--config=./partial"}, []string{"--config=./partial dir"}},
		{"short attached file", `probe -c./partial\ f`, pathReply(commandmeta.CompletionFiles, "-c"), []string{"__complete", "-c./partial f"}, []string{"-c./partial file"}},
		{"short cluster directory", "probe -sc./partial", pathReply(commandmeta.CompletionDirectories, "-sc"), []string{"__complete", "-sc./partial"}, []string{"-sc./partial dir"}},
		{"short path with equals", "probe -cfoo=ba", pathReply(commandmeta.CompletionFiles, "-c"), []string{"__complete", "-cfoo=ba"}, []string{"-cfoo=bar"}},
		{"short attached choice", "probe -cbl", choiceReply("-c", "blue"), []string{"__complete", "-cbl"}, []string{"-cblue"}},
		{"short equals choice", "probe -c=bl", choiceReply("-c=", "blue"), []string{"__complete", "-c=bl"}, []string{"-c=blue"}},
		{"short choice with equals", "probe -cfoo=ba", choiceReply("-c", "foo=bar"), []string{"__complete", "-cfoo=ba"}, []string{"-cfoo=bar"}},
		{"short no suggestions", "probe -cnomatch", pathReply(commandmeta.CompletionChoices, "-c"), []string{"__complete", "-cnomatch"}, []string{"-cnomatch"}},
		{"literal prefix marker", "probe __commandkit_path_pre", choiceReply("", completionPathPrefixMarker+"-c"), []string{"__complete", "__commandkit_path_pre"}, []string{completionPathPrefixMarker + "-c"}},
		{"literal marker attached", "probe --config=__commandkit_path_pre", choiceReply("--config=", completionPathPrefixMarker+"-c"), []string{"__complete", "--config=__commandkit_path_pre"}, []string{"--config=" + completionPathPrefixMarker + "-c"}},
		{"option-looking separate value", "probe --format --tag=ch", choiceReply("", "--tag=chosen"), []string{"__complete", "--format", "--tag=ch"}, []string{"--format", "--tag=chosen"}},
		{"option-looking operand after terminator", "probe -- --tag=ch", choiceReply("", "--tag=chosen"), []string{"__complete", "--", "--tag=ch"}, []string{"--", "--tag=chosen"}},
		{"empty", "probe ", reply("", spfcobra.ShellCompDirectiveNoFileComp), []string{"__complete", ""}, []string{}},
		{"no file fallback", `probe ./partial\ f`, reply("", spfcobra.ShellCompDirectiveNoFileComp), []string{"__complete", "./partial f"}, []string{"./partial f"}},
		{"error no fallback", `probe ./partial\ f`, reply("", spfcobra.ShellCompDirectiveError), []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", choiceReply("--config=", "./partial file"), []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)
				}
			})
		}
		if shell == "bash" || shell == "zsh" {
			for _, tc := range []struct{ name, line string }{
				{"quoted", `probe --format '-$(touch sentinel)=x'`},
				{"escaped", `probe --format -\$\(touch\ sentinel\)=x`},
			} {
				t.Run(shell+"/literal query/"+tc.name, func(t *testing.T) {
					log := filepath.Join(dir, "query-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.Env = append(os.Environ(), "PATH="+dir+string(os.PathListSeparator)+os.Getenv("PATH"), "HOME="+dir, "TERM=xterm", "PROBE_DIR="+dir, "PROBE_LOG="+log, "PROBE_TAB_ONLY=1", "PROBE_REPLY=:4\n")
					if out, err := cmd.CombinedOutput(); err != nil {
						t.Fatalf("tab-only completion: %v\n%s", err, out)
					}
					if _, err := os.Stat(filepath.Join(dir, "sentinel")); !os.IsNotExist(err) {
						t.Fatalf("completion executed literal input: %v", err)
					}
					got, err := os.ReadFile(log)
					if err != nil || !strings.Contains(string(got), `["__complete", "--format", "-$(touch sentinel)=x"]`) {
						t.Fatalf("literal query was not preserved: %q, %v", got, err)
					}
				})
			}
		}

		for _, executable := range []string{"probe.app", "café"} {
			t.Run(shell+"/executable "+executable, func(t *testing.T) {
				var generated bytes.Buffer
				if err := CompletionScript(&generated, executable, 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, executable+" --config=./par")
				cmd.Env = append(os.Environ(), "PATH="+dir+string(os.PathListSeparator)+os.Getenv("PATH"), "HOME="+dir, "TERM=xterm", "PROBE_LOG="+log, "PROBE_REPLY="+choiceReply("--config=", "./partial file"), "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
}