Fix misaligned command usage for 'help' commands of Teleport binaries (#51660)

* Fix an issue command help is not aligned for help command

* align when unknown command/subcommand

Author: greedy52

lib/utils/cli.go

@@ -391,25 +391,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()
 		}

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)
+			})
+		})
+	}
+}