Cancel context when command finishes

ammario · ammario · commit c6621bb7cdea · 2023-03-09T20:43:02.000Z
diff --git a/cli/clibase/cmd.go b/cli/clibase/cmd.go
@@ -251,6 +251,10 @@ func (i *Invokation) run(state *runState) error {
 		mw = Chain()
 	}
 
+	ctx, cancel := context.WithCancel(i.Context())
+	defer cancel()
+	i = i.WithContext(ctx)
+
 	if i.Command.Handler == nil || errors.Is(state.flagParseErr, pflag.ErrHelp) {
 		if i.Command.HelpHandler == nil {
 			return xerrors.Errorf("no handler or help for command %s", i.Command.FullName())
diff --git a/cli/clibase/cmd_test.go b/cli/clibase/cmd_test.go
@@ -1,6 +1,7 @@
 package clibase_test
 
 import (
+	"context"
 	"strings"
 	"testing"
 
@@ -327,3 +328,25 @@ func TestCommand_HyphenHypen(t *testing.T) {
 
 	require.Equal(t, "--verbose --friendly", stdio.Stdout.String())
 }
+
+func TestCommand_ContextCancels(t *testing.T) {
+	t.Parallel()
+
+	var gotCtx context.Context
+
+	cmd := &clibase.Cmd{
+		Handler: (func(i *clibase.Invokation) error {
+			gotCtx = i.Context()
+			if err := gotCtx.Err(); err != nil {
+				return xerrors.Errorf("unexpected context error: %w", i.Context().Err())
+			}
+			return nil
+		}),
+	}
+
+	inv, _ := clibasetest.Invoke(cmd)
+	err := inv.Run()
+	require.NoError(t, err)
+
+	require.Error(t, gotCtx.Err())
+}
