diff --git a/lib/utils/cli.go b/lib/utils/cli.go index a251d58ce41e3..faba9a2156d9c 100644 --- a/lib/utils/cli.go +++ b/lib/utils/cli.go @@ -380,25 +380,38 @@ func createUsageTemplate(opts ...func(*usageTemplateOptions)) string { // pre-parsing the arguments then applying any changes to the usage template if // necessary. func UpdateAppUsageTemplate(app *kingpin.Application, args []string) { - // If ParseContext fails, kingpin will not show usage so there is no need - // to update anything here. See app.Parse for more details. - context, err := app.ParseContext(args) - if err != nil { - return - } - app.UsageTemplate(createUsageTemplate( - withCommandPrintfWidth(app, context), + withCommandPrintfWidth(app, args), )) } -// withCommandPrintfWidth returns an usage template option that +// withCommandPrintfWidth returns a usage template option that // updates command printf width if longer than default. -func withCommandPrintfWidth(app *kingpin.Application, context *kingpin.ParseContext) func(*usageTemplateOptions) { +func withCommandPrintfWidth(app *kingpin.Application, args []string) func(*usageTemplateOptions) { return func(opt *usageTemplateOptions) { var commands []*kingpin.CmdModel - if context.SelectedCommand != nil { - commands = context.SelectedCommand.Model().FlattenedCommands() + + // When selected command is "help", skip the "help" arg + // so the intended command is selected for calculation. + if len(args) > 0 && args[0] == "help" { + args = args[1:] + } + + appContext, err := app.ParseContext(args) + switch { + case appContext == nil: + slog.WarnContext(context.Background(), "No application context found") + return + + // Note that ParseContext may return the current selected command that's + // causing the error. We should continue in those cases when appContext is + // not nil. + case err != nil: + slog.InfoContext(context.Background(), "Error parsing application context", "error", err) + } + + if appContext.SelectedCommand != nil { + commands = appContext.SelectedCommand.Model().FlattenedCommands() } else { commands = app.Model().FlattenedCommands() } diff --git a/lib/utils/cli_test.go b/lib/utils/cli_test.go index dcfccdeaab9cf..1c2e031c4d312 100644 --- a/lib/utils/cli_test.go +++ b/lib/utils/cli_test.go @@ -19,12 +19,14 @@ package utils import ( + "bytes" "crypto/x509" "fmt" "io" "log/slog" "testing" + "github.com/alecthomas/kingpin/v2" "github.com/gravitational/trace" "github.com/stretchr/testify/require" ) @@ -161,3 +163,82 @@ func TestAllowWhitespace(t *testing.T) { require.Equal(t, tt.out, AllowWhitespace(tt.in), fmt.Sprintf("test case %v", i)) } } + +func TestUpdateAppUsageTemplate(t *testing.T) { + makeApp := func(usageWriter io.Writer) *kingpin.Application { + app := InitCLIParser("TestUpdateAppUsageTemplate", "some help message") + app.UsageWriter(usageWriter) + app.Terminate(func(int) {}) + + app.Command("hello", "Hello.") + + create := app.Command("create", "Create.") + create.Command("box", "Box.") + create.Command("rocket", "Rocket.") + return app + } + + tests := []struct { + name string + inputArgs []string + outputContains string + }{ + { + name: "command width aligned for app help", + inputArgs: []string{}, + outputContains: ` +Commands: + help Show help. + hello Hello. + create box Box. + create rocket Rocket. +`, + }, + { + name: "command width aligned for command help", + inputArgs: []string{"create"}, + outputContains: ` +Commands: + create box Box. + create rocket Rocket. +`, + }, + { + name: "command width aligned for unknown command error", + inputArgs: []string{"unknown"}, + outputContains: ` +Commands: + help Show help. + hello Hello. + create box Box. + create rocket Rocket. +`, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Run("help flag", func(t *testing.T) { + var buffer bytes.Buffer + app := makeApp(&buffer) + args := append(tt.inputArgs, "--help") + UpdateAppUsageTemplate(app, args) + + app.Usage(args) + require.Contains(t, buffer.String(), tt.outputContains) + }) + + t.Run("help command", func(t *testing.T) { + var buffer bytes.Buffer + app := makeApp(&buffer) + args := append([]string{"help"}, tt.inputArgs...) + UpdateAppUsageTemplate(app, args) + + // HelpCommand is triggered on PreAction during Parse. + // See kingpin.Application.init for more details. + _, err := app.Parse(args) + require.NoError(t, err) + require.Contains(t, buffer.String(), tt.outputContains) + }) + }) + } +}