Repository: spf13/cobra
Pull request #2053 (2023-10-23 14:04:47)

Add tests for flag completion registration
Different problems have been reported about flag completion registration. 
These two tests are the cases that were not being verified but had been mentioned as problematic.

The test `TestFlagCompletionConcurrentRegistration` reproduces the `"concurrent map writes"` errors if concurrency is not properly handed, as reported by #1320.
I have verified this test does fail by running it on code before the fix from #1438.

The test `TestFlagCompletionForPersistentFlagsCalledFromSubCmd` makes sure persistent flags get completed.
This is based on problems mentions here:
- https://github.com/spf13/cobra/pull/1438#issuecomment-872928669
- https://github.com/spf13/cobra/pull/2012#pullrequestreview-1627336420

I have verified this test does fail based on the interim commit that was reported to have this issue in this [comment](https://github.com/spf13/cobra/pull/1438#issuecomment-872928669).


Patch:
diff --git a/command.go b/command.go
index ae3e4e0..36494df 100644
--- a/command.go
+++ b/command.go
@@ -1446,6 +1446,7 @@ func (c *Command) UseLine() string {
 
 // DebugFlags used to determine which flags have been assigned to which commands
 // and which persist.
+// nolint:goconst
 func (c *Command) DebugFlags() {
 	c.Println("DebugFlags called on", c.Name())
 	var debugflags func(*Command)
diff --git a/completions.go b/completions.go
index 55a9f62..368b92a 100644
--- a/completions.go
+++ b/completions.go
@@ -302,9 +302,13 @@ func (c *Command) getCompletions(args []string) (*Command, []string, ShellCompDi
 
 	// These flags are normally added when `execute()` is called on `finalCmd`,
 	// however, when doing completion, we don't call `finalCmd.execute()`.
-	// Let's add the --help and --version flag ourselves.
-	finalCmd.InitDefaultHelpFlag()
-	finalCmd.InitDefaultVersionFlag()
+	// Let's add the --help and --version flag ourselves but only if the finalCmd
+	// has not disabled flag parsing; if flag parsing is disabled, it is up to the
+	// finalCmd itself to handle the completion of *all* flags.
+	if !finalCmd.DisableFlagParsing {
+		finalCmd.InitDefaultHelpFlag()
+		finalCmd.InitDefaultVersionFlag()
+	}
 
 	// Check if we are doing flag value completion before parsing the flags.
 	// This is important because if we are completing a flag value, we need to also
@@ -408,6 +412,11 @@ func (c *Command) getCompletions(args []string) (*Command, []string, ShellCompDi
 			finalCmd.InheritedFlags().VisitAll(func(flag *pflag.Flag) {
 				doCompleteFlags(flag)
 			})
+			// Try to complete non-inherited flags even if DisableFlagParsing==true.
+			// This allows programs to tell Cobra about flags for completion even
+			// if the actual parsing of flags is not done by Cobra.
+			// For instance, Helm uses this to provide flag name completion for
+			// some of its plugins.
 			finalCmd.NonInheritedFlags().VisitAll(func(flag *pflag.Flag) {
 				doCompleteFlags(flag)
 			})
diff --git a/doc/md_docs.go b/doc/md_docs.go
index c4a27c0..f98fe2a 100644
--- a/doc/md_docs.go
+++ b/doc/md_docs.go
@@ -27,6 +27,8 @@ import (
 	"github.com/spf13/cobra"
 )
 
+const markdownExtension = ".md"
+
 func printOptions(buf *bytes.Buffer, cmd *cobra.Command, name string) error {
 	flags := cmd.NonInheritedFlags()
 	flags.SetOutput(buf)
@@ -83,7 +85,7 @@ func GenMarkdownCustom(cmd *cobra.Command, w io.Writer, linkHandler func(string)
 		if cmd.HasParent() {
 			parent := cmd.Parent()
 			pname := parent.CommandPath()
-			link := pname + ".md"
+			link := pname + markdownExtension
 			link = strings.ReplaceAll(link, " ", "_")
 			buf.WriteString(fmt.Sprintf("* [%s](%s)\t - %s\n", pname, linkHandler(link), parent.Short))
 			cmd.VisitParents(func(c *cobra.Command) {
@@ -101,7 +103,7 @@ func GenMarkdownCustom(cmd *cobra.Command, w io.Writer, linkHandler func(string)
 				continue
 			}
 			cname := name + " " + child.Name()
-			link := cname + ".md"
+			link := cname + markdownExtension
 			link = strings.ReplaceAll(link, " ", "_")
 			buf.WriteString(fmt.Sprintf("* [%s](%s)\t - %s\n", cname, linkHandler(link), child.Short))
 		}
@@ -138,7 +140,7 @@ func GenMarkdownTreeCustom(cmd *cobra.Command, dir string, filePrepender, linkHa
 		}
 	}
 
-	basename := strings.ReplaceAll(cmd.CommandPath(), " ", "_") + ".md"
+	basename := strings.ReplaceAll(cmd.CommandPath(), " ", "_") + markdownExtension
 	filename := filepath.Join(dir, basename)
 	f, err := os.Create(filename)
 	if err != nil {


Test patch:
diff --git a/completions_test.go b/completions_test.go
index 0172466..7585f88 100644
--- a/completions_test.go
+++ b/completions_test.go
@@ -17,7 +17,9 @@ package cobra
 import (
 	"bytes"
 	"context"
+	"fmt"
 	"strings"
+	"sync"
 	"testing"
 )
 
@@ -2040,6 +2042,114 @@ func TestFlagCompletionWorksRootCommandAddedAfterFlags(t *testing.T) {
 	}
 }
 
+func TestFlagCompletionForPersistentFlagsCalledFromSubCmd(t *testing.T) {
+	rootCmd := &Command{Use: "root", Run: emptyRun}
+	rootCmd.PersistentFlags().String("string", "", "test string flag")
+	_ = rootCmd.RegisterFlagCompletionFunc("string", func(cmd *Command, args []string, toComplete string) ([]string, ShellCompDirective) {
+		return []string{"myval"}, ShellCompDirectiveDefault
+	})
+
+	childCmd := &Command{
+		Use: "child",
+		Run: emptyRun,
+		ValidArgsFunction: func(cmd *Command, args []string, toComplete string) ([]string, ShellCompDirective) {
+			return []string{"--validarg", "test"}, ShellCompDirectiveDefault
+		},
+	}
+	childCmd.Flags().Bool("bool", false, "test bool flag")
+	rootCmd.AddCommand(childCmd)
+
+	// Test that persistent flag completion works for the subcmd
+	output, err := executeCommand(rootCmd, ShellCompRequestCmd, "child", "--string", "")
+	if err != nil {
+		t.Errorf("Unexpected error: %v", err)
+	}
+
+	expected := strings.Join([]string{
+		"myval",
+		":0",
+		"Completion ended with directive: ShellCompDirectiveDefault", ""}, "\n")
+
+	if output != expected {
+		t.Errorf("expected: %q, got: %q", expected, output)
+	}
+}
+
+// This test tries to register flag completion concurrently to make sure the
+// code handles concurrency properly.
+// This was reported as a problem when tests are run concurrently:
+// https://github.com/spf13/cobra/issues/1320
+//
+// NOTE: this test can sometimes pass even if the code were to not handle
+// concurrency properly. This is not great but the important part is that
+// it should never fail.  Therefore, if the tests fails sometimes, we will
+// still be able to know there is a problem.
+func TestFlagCompletionConcurrentRegistration(t *testing.T) {
+	rootCmd := &Command{Use: "root", Run: emptyRun}
+	const maxFlags = 50
+	for i := 1; i < maxFlags; i += 2 {
+		flagName := fmt.Sprintf("flag%d", i)
+		rootCmd.Flags().String(flagName, "", fmt.Sprintf("test %s flag on root", flagName))
+	}
+
+	childCmd := &Command{
+		Use: "child",
+		Run: emptyRun,
+	}
+	for i := 2; i <= maxFlags; i += 2 {
+		flagName := fmt.Sprintf("flag%d", i)
+		childCmd.Flags().String(flagName, "", fmt.Sprintf("test %s flag on child", flagName))
+	}
+
+	rootCmd.AddCommand(childCmd)
+
+	// Register completion in different threads to test concurrency.
+	var wg sync.WaitGroup
+	for i := 1; i <= maxFlags; i++ {
+		index := i
+		flagName := fmt.Sprintf("flag%d", i)
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			cmd := rootCmd
+			if index%2 == 0 {
+				cmd = childCmd
+			}
+			_ = cmd.RegisterFlagCompletionFunc(flagName, func(cmd *Command, args []string, toComplete string) ([]string, ShellCompDirective) {
+				return []string{fmt.Sprintf("flag%d", index)}, ShellCompDirectiveDefault
+			})
+		}()
+	}
+
+	wg.Wait()
+
+	// Test that flag completion works for each flag
+	for i := 1; i <= 6; i++ {
+		var output string
+		var err error
+		flagName := fmt.Sprintf("flag%d", i)
+
+		if i%2 == 1 {
+			output, err = executeCommand(rootCmd, ShellCompRequestCmd, "--"+flagName, "")
+		} else {
+			output, err = executeCommand(rootCmd, ShellCompRequestCmd, "child", "--"+flagName, "")
+		}
+
+		if err != nil {
+			t.Errorf("Unexpected error: %v", err)
+		}
+
+		expected := strings.Join([]string{
+			flagName,
+			":0",
+			"Completion ended with directive: ShellCompDirectiveDefault", ""}, "\n")
+
+		if output != expected {
+			t.Errorf("expected: %q, got: %q", expected, output)
+		}
+	}
+}
+
 func TestFlagCompletionInGoWithDesc(t *testing.T) {
 	rootCmd := &Command{
 		Use: "root",
@@ -2622,8 +2732,6 @@ func TestCompleteWithDisableFlagParsing(t *testing.T) {
 	expected := strings.Join([]string{
 		"--persistent",
 		"-p",
-		"--help",
-		"-h",
 		"--nonPersistent",
 		"-n",
 		"--flag",
@@ -3053,8 +3161,26 @@ func TestCompletionCobraFlags(t *testing.T) {
 				return []string{"extra3"}, ShellCompDirectiveNoFileComp
 			},
 		}
+		childCmd4 := &Command{
+			Use:     "child4",
+			Version: "1.1.1",
+			Run:     emptyRun,
+			ValidArgsFunction: func(cmd *Command, args []string, toComplete string) ([]string, ShellCompDirective) {
+				return []string{"extra4"}, ShellCompDirectiveNoFileComp
+			},
+			DisableFlagParsing: true,
+		}
+		childCmd5 := &Command{
+			Use:     "child5",
+			Version: "1.1.1",
+			Run:     emptyRun,
+			ValidArgsFunction: func(cmd *Command, args []string, toComplete string) ([]string, ShellCompDirective) {
+				return []string{"extra5"}, ShellCompDirectiveNoFileComp
+			},
+			DisableFlagParsing: true,
+		}
 
-		rootCmd.AddCommand(childCmd, childCmd2, childCmd3)
+		rootCmd.AddCommand(childCmd, childCmd2, childCmd3, childCmd4, childCmd5)
 
 		_ = childCmd.Flags().Bool("bool", false, "A bool flag")
 		_ = childCmd.MarkFlagRequired("bool")
@@ -3066,6 +3192,10 @@ func TestCompletionCobraFlags(t *testing.T) {
 		// Have a command that only adds its own -v flag
 		_ = childCmd3.Flags().BoolP("verbose", "v", false, "Not a version flag")
 
+		// Have a command that DisablesFlagParsing but that also adds its own help and version flags
+		_ = childCmd5.Flags().BoolP("help", "h", false, "My own help")
+		_ = childCmd5.Flags().BoolP("version", "v", false, "My own version")
+
 		return rootCmd
 	}
 
@@ -3196,6 +3326,26 @@ func TestCompletionCobraFlags(t *testing.T) {
 				":4",
 				"Completion ended with directive: ShellCompDirectiveNoFileComp", ""}, "\n"),
 		},
+		{
+			desc: "no completion for --help/-h and --version/-v flags when DisableFlagParsing=true",
+			args: []string{"child4", "-"},
+			expectedOutput: strings.Join([]string{
+				"extra4",
+				":4",
+				"Completion ended with directive: ShellCompDirectiveNoFileComp", ""}, "\n"),
+		},
+		{
+			desc: "completions for program-defined --help/-h and --version/-v flags even when DisableFlagParsing=true",
+			args: []string{"child5", "-"},
+			expectedOutput: strings.Join([]string{
+				"--help",
+				"-h",
+				"--version",
+				"-v",
+				"extra5",
+				":4",
+				"Completion ended with directive: ShellCompDirectiveNoFileComp", ""}, "\n"),
+		},
 	}
 
 	for _, tc := range testcases {