diff --git a/ext/ext.go b/ext/ext.go index e1638391ff..6917a42807 100644 --- a/ext/ext.go +++ b/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 } @@ -103,8 +106,8 @@ func GetAll() []*Extension { mx.RLock() defer mx.RUnlock() - js, out := extensions[JSExtension], extensions[OutputExtension] - result := make([]*Extension, 0, len(js)+len(out)) + js, out, subcommand := extensions[JSExtension], extensions[OutputExtension], extensions[SubcommandExtension] + result := make([]*Extension, 0, len(js)+len(out)+len(subcommand)) for _, e := range js { result = append(result, e) @@ -112,6 +115,9 @@ func GetAll() []*Extension { for _, e := range out { result = append(result, e) } + for _, e := range subcommand { + result = append(result, e) + } sort.Slice(result, func(i, j int) bool { if result[i].Path == result[j].Path { @@ -161,4 +167,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) } diff --git a/internal/cmd/root.go b/internal/cmd/root.go index 7fcc02d2ac..b6a3cf5e01 100644 --- a/internal/cmd/root.go +++ b/internal/cmd/root.go @@ -87,6 +87,17 @@ func newRootCommand(gs *state.GlobalState) *rootCommand { rootCmd.AddCommand(sc(gs)) } + xCmd := getX(gs) + + // Add extension subcommands + for sc := range extensionSubcommands(gs, xCmd.Commands()) { + xCmd.AddCommand(sc) + } + + if len(xCmd.Commands()) > 0 { + rootCmd.AddCommand(xCmd) + } + c.cmd = rootCmd return c } diff --git a/internal/cmd/root_test.go b/internal/cmd/root_test.go index cb3a12056c..a2e9e72193 100644 --- a/internal/cmd/root_test.go +++ b/internal/cmd/root_test.go @@ -11,6 +11,8 @@ import ( func TestRootCommandHelpDisplayCommands(t *testing.T) { t.Parallel() + registerTestSubcommandExtensions(t) + testCases := []struct { name string extraArgs []string @@ -70,6 +72,10 @@ func TestRootCommandHelpDisplayCommands(t *testing.T) { name: "should have version command", wantStdoutContains: " version Show application version", }, + { + name: "should have x command", + wantStdoutContains: " x Extension subcommands", + }, } for _, tc := range testCases { diff --git a/internal/cmd/subcommand.go b/internal/cmd/subcommand.go new file mode 100644 index 0000000000..cfb25e0f76 --- /dev/null +++ b/internal/cmd/subcommand.go @@ -0,0 +1,66 @@ +package cmd + +import ( + "fmt" + "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" +) + +func getX(_ *state.GlobalState) *cobra.Command { + return &cobra.Command{ + Use: "x", + Short: "Extension subcommands", + Long: `Namespace for extension-provided subcommands. + +This command serves as a parent for subcommands registered by k6 extensions, +allowing them to extend k6's functionality with custom commands. +`, + } +} + +// 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 { + panic(fmt.Sprintf("invalid subcommand constructor: name: %s path: %s", extension.Name, extension.Path)) + } + + cmd := ctor(gs) + + // Validate that the command's name matches the extension name. + if cmd.Name() != extension.Name { + panic(fmt.Sprintf("subcommand name mismatch: command name: %s extension name: %s", cmd.Name(), extension.Name)) + } + + return cmd +} diff --git a/internal/cmd/subcommand_test.go b/internal/cmd/subcommand_test.go new file mode 100644 index 0000000000..2ccfba23b1 --- /dev/null +++ b/internal/cmd/subcommand_test.go @@ -0,0 +1,185 @@ +package cmd + +import ( + "sync" + "testing" + + "github.com/spf13/cobra" + "github.com/stretchr/testify/require" + "go.k6.io/k6/cmd/state" + "go.k6.io/k6/internal/cmd/tests" + "go.k6.io/k6/subcommand" +) + +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) + } + } + }) +} + +func TestXCommandHelpDisplayCommands(t *testing.T) { + t.Parallel() + + registerTestSubcommandExtensions(t) + + testCases := []struct { + name string + wantStdoutContains string + }{ + { + 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 { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + ts := tests.NewGlobalTestState(t) + ts.CmdArgs = []string{"k6", "x", "help"} + newRootCommand(ts.GlobalState).execute() + + require.Contains(t, ts.Stdout.String(), tc.wantStdoutContains) + }) + } +} + +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) {}, + } + }) + }) +} diff --git a/subcommand/extension.go b/subcommand/extension.go new file mode 100644 index 0000000000..b80d42ada3 --- /dev/null +++ b/subcommand/extension.go @@ -0,0 +1,36 @@ +// 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. +// +// WARNING: The GlobalState parameter is read-only and must not be modified or altered +// in any way. Modifying the GlobalState can make k6 core unstable and lead to +// unpredictable behavior. +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 "). +// 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) +}