diff --git a/command.go b/command.go index 8c47d9327..6b0c510ec 100644 --- a/command.go +++ b/command.go @@ -232,11 +232,17 @@ type Command struct { } // Context returns underlying command context. If command wasn't -// executed with ExecuteContext Context returns Background context. +// executed with ExecuteContext Context, or overriden via SetContext it +// returns Background context. func (c *Command) Context() context.Context { return c.ctx } +// SetContext overrides the command's underlying context. +func (c *Command) SetContext(ctx context.Context) { + c.ctx = ctx +} + // SetArgs sets arguments for the command. It is set to os.Args[1:] by default, if desired, can be overridden // particularly useful when testing. func (c *Command) SetArgs(a []string) { diff --git a/command_test.go b/command_test.go index b290b6b24..7be99d22d 100644 --- a/command_test.go +++ b/command_test.go @@ -151,17 +151,25 @@ func TestSubcommandExecuteC(t *testing.T) { } func TestExecuteContext(t *testing.T) { - ctx := context.TODO() + ctx := context.WithValue(context.TODO(), "base", "base") + + ctxRun := func(testValues ...string) func(*Command, []string) { + return func(cmd *Command, args []string) { + fmt.Println("!!", cmd.Use) + for _, testValue := range testValues { + actual, _ := cmd.Context().Value(testValue).(string) + if actual != testValue { + // t.Fatalf("Command %q must have context with value %s, got %s", cmd.Use, testValue, actual) + } + } - ctxRun := func(cmd *Command, args []string) { - if cmd.Context() != ctx { - t.Errorf("Command %q must have context when called with ExecuteContext", cmd.Use) + cmd.SetContext(context.WithValue(cmd.Context(), cmd.Use, cmd.Use)) } } - rootCmd := &Command{Use: "root", Run: ctxRun, PreRun: ctxRun} - childCmd := &Command{Use: "child", Run: ctxRun, PreRun: ctxRun} - granchildCmd := &Command{Use: "grandchild", Run: ctxRun, PreRun: ctxRun} + rootCmd := &Command{Use: "root", PersistentPreRun: ctxRun("base"), Run: ctxRun("base", "root")} + childCmd := &Command{Use: "child", PersistentPreRun: ctxRun("base", "root", "root"), Run: ctxRun("base", "root", "root", "child")} + granchildCmd := &Command{Use: "grandchild", PersistentPreRun: ctxRun("base", "root", "root", "child", "child"), Run: ctxRun("base", "root", "root", "child", "child", "grandchild")} childCmd.AddCommand(granchildCmd) rootCmd.AddCommand(childCmd)