Add subcommand extension support

ext/ext.go

@@ -26,6 +26,7 @@ const (
 	JSExtension ExtensionType = iota + 1
 	OutputExtension
 	SecretSourceExtension
+	SubcommandExtension
 )
 
 func (e ExtensionType) String() string {
@@ -37,6 +38,8 @@ func (e ExtensionType) String() string {
 		s = "output"
 	case SecretSourceExtension:
 		s = "secret-source"
+	case SubcommandExtension:
+		s = "subcommand"
 	}
 	return s
 }
@@ -161,4 +164,5 @@ func init() {
 	extensions[JSExtension] = make(map[string]*Extension)
 	extensions[OutputExtension] = make(map[string]*Extension)
 	extensions[SecretSourceExtension] = make(map[string]*Extension)
+	extensions[SubcommandExtension] = make(map[string]*Extension)
 }

internal/cmd/root.go

@@ -87,6 +87,10 @@ func newRootCommand(gs *state.GlobalState) *rootCommand {
 		rootCmd.AddCommand(sc(gs))
 	}
 
+	for sc := range extensionSubcommands(gs, rootCmd.Commands()) {
+		rootCmd.AddCommand(sc)
+	}
+
 	c.cmd = rootCmd
 	return c
 }

internal/cmd/root_test.go

@@ -1,16 +1,22 @@
 package cmd
 
 import (
+	"sync"
 	"testing"
 
+	"github.com/spf13/cobra"
 	"github.com/stretchr/testify/assert"
+	"go.k6.io/k6/cmd/state"
 	"go.k6.io/k6/errext/exitcodes"
 	"go.k6.io/k6/internal/cmd/tests"
+	"go.k6.io/k6/subcommand"
 )
 
 func TestRootCommandHelpDisplayCommands(t *testing.T) {
 	t.Parallel()
 
+	registerTestSubcommandExtensions(t)
+
 	testCases := []struct {
 		name                  string
 		extraArgs             []string
@@ -70,6 +76,18 @@ func TestRootCommandHelpDisplayCommands(t *testing.T) {
 			name:               "should have version command",
 			wantStdoutContains: "  version     Show application version",
 		},
+		{
+			name:               "should have test-cmd-1 command",
+			wantStdoutContains: "  test-cmd-1  Test command 1",
+		},
+		{
+			name:               "should have test-cmd-2 command",
+			wantStdoutContains: "  test-cmd-2  Test command 2",
+		},
+		{
+			name:               "should have test-cmd-3 command",
+			wantStdoutContains: "  test-cmd-3  Test command 3",
+		},
 	}
 
 	for _, tc := range testCases {
@@ -86,3 +104,35 @@ func TestRootCommandHelpDisplayCommands(t *testing.T) {
 		})
 	}
 }
+
+var registerTestSubcommandExtensionsOnce sync.Once //nolint:gochecknoglobals
+
+func registerTestSubcommandExtensions(t *testing.T) {
+	t.Helper()
+
+	registerTestSubcommandExtensionsOnce.Do(func() {
+		subcommand.RegisterExtension("test-cmd-1", func(_ *state.GlobalState) *cobra.Command {
+			return &cobra.Command{
+				Use:   "test-cmd-1",
+				Short: "Test command 1",
+				Run:   func(_ *cobra.Command, _ []string) {},
+			}
+		})
+
+		subcommand.RegisterExtension("test-cmd-2", func(_ *state.GlobalState) *cobra.Command {
+			return &cobra.Command{
+				Use:   "test-cmd-2",
+				Short: "Test command 2",
+				Run:   func(_ *cobra.Command, _ []string) {},
+			}
+		})
+
+		subcommand.RegisterExtension("test-cmd-3", func(_ *state.GlobalState) *cobra.Command {
+			return &cobra.Command{
+				Use:   "test-cmd-3",
+				Short: "Test command 3",
+				Run:   func(_ *cobra.Command, _ []string) {},
+			}
+		})
+	})
+}

internal/cmd/subcommand.go

@@ -0,0 +1,55 @@
+package cmd
+
+import (
+	"iter"
+
+	"github.com/sirupsen/logrus"
+	"github.com/spf13/cobra"
+	"go.k6.io/k6/cmd/state"
+	"go.k6.io/k6/ext"
+	"go.k6.io/k6/subcommand"
+)
+
+// extensionSubcommands returns an iterator over all registered subcommand extensions
+// that are not already defined in the given slice of commands.
+func extensionSubcommands(gs *state.GlobalState, defined []*cobra.Command) iter.Seq[*cobra.Command] {
+	already := make(map[string]struct{}, len(defined))
+	for _, cmd := range defined {
+		already[cmd.Name()] = struct{}{}
+	}
+
+	return func(yield func(*cobra.Command) bool) {
+		for _, extension := range ext.Get(ext.SubcommandExtension) {
+			if _, exists := already[extension.Name]; exists {
+				gs.Logger.WithFields(logrus.Fields{"name": extension.Name, "path": extension.Path}).
+					Warnf("subcommand already exists")
+				continue
+			}
+
+			already[extension.Name] = struct{}{}
+
+			if !yield(getCmdForExtension(extension, gs)) {
+				break
+			}
+		}
+	}
+}
+
+// getCmdForExtension gets a *cobra.Command for the given subcommand extension.
+func getCmdForExtension(extension *ext.Extension, gs *state.GlobalState) *cobra.Command {
+	ctor, ok := extension.Module.(subcommand.Constructor)
+	if !ok {
+		gs.Logger.WithFields(logrus.Fields{"name": extension.Name, "path": extension.Path}).
+			Fatalf("subcommand's constructor does not implement the subcommand.Constructor")
+	}
+
+	cmd := ctor(gs)
+
+	// Validate that the command's name matches the extension name.
+	if cmd.Name() != extension.Name {
+		gs.Logger.WithFields(logrus.Fields{"name": extension.Name, "path": extension.Path}).
+			Fatalf("subcommand's command name (%s) does not match the extension name (%s)", cmd.Name(), extension.Name)
+	}
+
+	return cmd
+}

internal/cmd/subcommand_test.go

@@ -0,0 +1,114 @@
+package cmd
+
+import (
+	"testing"
+
+	"github.com/spf13/cobra"
+	"github.com/stretchr/testify/require"
+	"go.k6.io/k6/internal/cmd/tests"
+)
+
+func TestExtensionSubcommands(t *testing.T) {
+	t.Parallel()
+
+	registerTestSubcommandExtensions(t)
+
+	t.Run("returns all extension subcommands", func(t *testing.T) {
+		t.Parallel()
+
+		ts := tests.NewGlobalTestState(t)
+		defined := []*cobra.Command{}
+
+		var commands []*cobra.Command
+		for cmd := range extensionSubcommands(ts.GlobalState, defined) {
+			commands = append(commands, cmd)
+		}
+
+		// Should have at least the 3 test extensions we registered
+		require.GreaterOrEqual(t, len(commands), 2)
+
+		// Check that our test commands are present
+		commandNames := make(map[string]bool)
+		for _, cmd := range commands {
+			commandNames[cmd.Name()] = true
+		}
+
+		require.True(t, commandNames["test-cmd-1"], "test-cmd-1 should be present")
+		require.True(t, commandNames["test-cmd-2"], "test-cmd-2 should be present")
+		require.True(t, commandNames["test-cmd-3"], "test-cmd-3 should be present")
+	})
+
+	t.Run("filters out already defined commands", func(t *testing.T) {
+		t.Parallel()
+
+		ts := tests.NewGlobalTestState(t)
+
+		// Create a command with the same name as one of our extensions
+		defined := []*cobra.Command{
+			{
+				Use:   "test-cmd-1",
+				Short: "Already defined command",
+				Run:   func(_ *cobra.Command, _ []string) {},
+			},
+		}
+
+		var commands []*cobra.Command
+		for cmd := range extensionSubcommands(ts.GlobalState, defined) {
+			commands = append(commands, cmd)
+		}
+
+		// Check that test-cmd-1 is NOT in the results
+		for _, cmd := range commands {
+			require.NotEqual(t, "test-cmd-1", cmd.Name(), "test-cmd-1 should be filtered out")
+		}
+
+		// But test-cmd-2 and test-cmd-3 should still be present
+		commandNames := make(map[string]bool)
+		for _, cmd := range commands {
+			commandNames[cmd.Name()] = true
+		}
+
+		require.True(t, commandNames["test-cmd-2"], "test-cmd-2 should be present")
+		require.True(t, commandNames["test-cmd-3"], "test-cmd-3 should be present")
+	})
+
+	t.Run("prevents duplicate extensions", func(t *testing.T) {
+		t.Parallel()
+
+		ts := tests.NewGlobalTestState(t)
+		defined := []*cobra.Command{}
+
+		// Collect all commands
+		var commands []*cobra.Command
+		for cmd := range extensionSubcommands(ts.GlobalState, defined) {
+			commands = append(commands, cmd)
+		}
+
+		// Check for duplicates
+		seen := make(map[string]bool)
+		for _, cmd := range commands {
+			require.False(t, seen[cmd.Name()], "command %s should not appear twice", cmd.Name())
+			seen[cmd.Name()] = true
+		}
+	})
+
+	t.Run("returns commands with correct properties", func(t *testing.T) {
+		t.Parallel()
+
+		ts := tests.NewGlobalTestState(t)
+		defined := []*cobra.Command{}
+
+		for cmd := range extensionSubcommands(ts.GlobalState, defined) {
+			require.NotEmpty(t, cmd.Use, "command should have a Use field")
+
+			switch cmd.Use {
+			case "test-cmd-1":
+				require.Equal(t, "Test command 1", cmd.Short)
+			case "test-cmd-2":
+				require.Equal(t, "Test command 2", cmd.Short)
+			case "test-cmd-3":
+				require.Equal(t, "Test command 3", cmd.Short)
+			}
+		}
+	})
+}

subcommand/extension.go

@@ -0,0 +1,32 @@
+// Package subcommand provides functionality for registering k6 subcommand extensions.
+//
+// This package allows external modules to register new subcommands that will be
+// available in the k6 CLI. Subcommand extensions are registered during
+// package initialization and are called when the corresponding subcommand is invoked.
+package subcommand
+
+import (
+	"github.com/spf13/cobra"
+	"go.k6.io/k6/cmd/state"
+	"go.k6.io/k6/ext"
+)
+
+// Constructor is a function type that creates a new cobra.Command for a subcommand extension.
+// It receives a GlobalState instance that provides access to configuration, logging,
+// file system, and other shared k6 runtime state. The returned Command will be
+// integrated into k6's CLI as a subcommand.
+type Constructor func(*state.GlobalState) *cobra.Command
+
+// RegisterExtension registers a subcommand extension with the given name and constructor function.
+//
+// The name parameter specifies the subcommand name that users will invoke (e.g., "k6 <name>").
+// The constructor function will be called when k6 initializes to create the cobra.Command
+// instance for this subcommand.
+//
+// This function must be called during package initialization (typically in an init() function)
+// and will panic if a subcommand with the same name is already registered.
+//
+// The name parameter and the returned Command's Name() must match.
+func RegisterExtension(name string, c Constructor) {
+	ext.Register(name, ext.SubcommandExtension, c)
+}