Pass context to completion (#1265)

completions_test.go

@@ -2,6 +2,7 @@ package cobra
 
 import (
 	"bytes"
+	"context"
 	"strings"
 	"testing"
 )
@@ -1203,6 +1204,48 @@ func TestFlagDirFilterCompletionInGo(t *testing.T) {
 	}
 }
 
+func TestValidArgsFuncCmdContext(t *testing.T) {
+	validArgsFunc := func(cmd *Command, args []string, toComplete string) ([]string, ShellCompDirective) {
+		ctx := cmd.Context()
+
+		if ctx == nil {
+			t.Error("Received nil context in completion func")
+		} else if ctx.Value("testKey") != "123" {
+			t.Error("Received invalid context")
+		}
+
+		return nil, ShellCompDirectiveDefault
+	}
+
+	rootCmd := &Command{
+		Use: "root",
+		Run: emptyRun,
+	}
+	childCmd := &Command{
+		Use:               "childCmd",
+		ValidArgsFunction: validArgsFunc,
+		Run:               emptyRun,
+	}
+	rootCmd.AddCommand(childCmd)
+
+	//nolint:golint,staticcheck // We can safely use a basic type as key in tests.
+	ctx := context.WithValue(context.Background(), "testKey", "123")
+
+	// Test completing an empty string on the childCmd
+	_, output, err := executeCommandWithContextC(ctx, rootCmd, ShellCompNoDescRequestCmd, "childCmd", "")
+	if err != nil {
+		t.Errorf("Unexpected error: %v", err)
+	}
+
+	expected := strings.Join([]string{
+		":0",
+		"Completion ended with directive: ShellCompDirectiveDefault", ""}, "\n")
+
+	if output != expected {
+		t.Errorf("expected: %q, got: %q", expected, output)
+	}
+}
+
 func TestValidArgsFuncSingleCmd(t *testing.T) {
 	rootCmd := &Command{
 		Use:               "root",