[refactor][kubectl-plugin] use cobra's posargs length validator

We're currently validating the number of positional arguments ourselves. We can do this with `cobra.Command{Args: cobra.MaximumNArgs(1)}` instead to simplify logic. There are no functional changes. The error message changes slightly but still conveys the same information. ## Before ```console $ kubectl ray --context gke_kubeflow-platform_europe-west4-b_ml-compute-1 get cluster foo bar Error: too many arguments, either one or no arguments are allowed ``` ## After ```console $ kubectl ray get cluster foo bar Error: accepts at most 1 arg(s), received 2 ``` Signed-off-by: David Xia <[email protected]>
ray-project · Feb 8, 2025 · 1cf6d45 · 1cf6d45
1 parent ffa6a30
commit 1cf6d45
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 38 deletions.
diff --git a/kubectl-plugin/pkg/cmd/get/get_cluster.go b/kubectl-plugin/pkg/cmd/get/get_cluster.go
@@ -22,7 +22,7 @@ import (
 type GetClusterOptions struct {
 	configFlags   *genericclioptions.ConfigFlags
 	ioStreams     *genericclioptions.IOStreams
-	args          []string
+	cluster       string
 	AllNamespaces bool
 }
 
@@ -43,6 +43,7 @@ func NewGetClusterCommand(streams genericclioptions.IOStreams) *cobra.Command {
 		Short:             "Get cluster information.",
 		SilenceUsage:      true,
 		ValidArgsFunction: completion.RayClusterCompletionFunc(cmdFactory),
+		Args:              cobra.MaximumNArgs(1),
 		RunE: func(cmd *cobra.Command, args []string) error {
 			if err := options.Complete(args); err != nil {
 				return err
@@ -68,7 +69,10 @@ func (options *GetClusterOptions) Complete(args []string) error {
 		options.AllNamespaces = true
 	}
 
-	options.args = args
+	if len(args) >= 1 {
+		options.cluster = args[0]
+	}
+
 	return nil
 }
 
@@ -81,9 +85,6 @@ func (options *GetClusterOptions) Validate() error {
 	if !util.HasKubectlContext(config, options.configFlags) {
 		return fmt.Errorf("no context is currently set, use %q or %q to select a new one", "--context", "kubectl config use-context <context>")
 	}
-	if len(options.args) > 1 {
-		return fmt.Errorf("too many arguments, either one or no arguments are allowed")
-	}
 	return nil
 }
 
@@ -92,21 +93,21 @@ func (options *GetClusterOptions) Run(ctx context.Context, k8sClient client.Clie
 	var rayclusterList *rayv1.RayClusterList
 
 	listopts := v1.ListOptions{}
-	if len(options.args) == 1 {
+	if options.cluster != "" {
 		listopts = v1.ListOptions{
-			FieldSelector: fmt.Sprintf("metadata.name=%s", options.args[0]),
+			FieldSelector: fmt.Sprintf("metadata.name=%s", options.cluster),
 		}
 	}
 
 	if options.AllNamespaces {
 		rayclusterList, err = k8sClient.RayClient().RayV1().RayClusters("").List(ctx, listopts)
 		if err != nil {
-			return fmt.Errorf("unable to retrieve raycluster for all namespaces: %w", err)
+			return fmt.Errorf("unable to retrieve Ray clusters for all namespaces: %w", err)
 		}
 	} else {
 		rayclusterList, err = k8sClient.RayClient().RayV1().RayClusters(*options.configFlags.Namespace).List(ctx, listopts)
 		if err != nil {
-			return fmt.Errorf("unable to retrieve raycluster for namespace %s: %w", *options.configFlags.Namespace, err)
+			return fmt.Errorf("unable to retrieve Ray clusters for namespace %s: %w", *options.configFlags.Namespace, err)
 		}
 	}
 

diff --git a/kubectl-plugin/pkg/cmd/get/get_cluster_test.go b/kubectl-plugin/pkg/cmd/get/get_cluster_test.go
@@ -21,22 +21,62 @@ import (
 	rayClientFake "github.com/ray-project/kuberay/ray-operator/pkg/client/clientset/versioned/fake"
 )
 
-// This is to test Complete() and ensure that it is setting the namespace and arguments correctly
+// This is to test Complete() and ensure that it is setting the namespace and cluster correctly
 // No validation test is done here
 func TestRayClusterGetComplete(t *testing.T) {
 	// Initialize members of the cluster get option struct and the struct itself
 	testStreams, _, _, _ := genericclioptions.NewTestIOStreams()
-	fakeClusterGetOptions := NewGetClusterOptions(testStreams)
-	fakeArgs := []string{"Expected", "output"}
 
-	*fakeClusterGetOptions.configFlags.Namespace = ""
-	fakeClusterGetOptions.AllNamespaces = false
+	tests := []struct {
+		name                  string
+		namespace             string
+		expectedCluster       string
+		args                  []string
+		expectedAllNamespaces bool
+	}{
+		{
+			name:                  "neither namespace nor args set",
+			namespace:             "",
+			args:                  []string{},
+			expectedAllNamespaces: true,
+			expectedCluster:       "",
+		},
+		{
+			name:                  "namespace set, args not set",
+			namespace:             "foo",
+			args:                  []string{},
+			expectedAllNamespaces: false,
+			expectedCluster:       "",
+		},
+		{
+			name:                  "namespace not set, args set",
+			namespace:             "",
+			args:                  []string{"foo", "bar"},
+			expectedAllNamespaces: true,
+			expectedCluster:       "foo",
+		},
+		{
+			name:                  "both namespace and args set",
+			namespace:             "foo",
+			args:                  []string{"bar", "qux"},
+			expectedAllNamespaces: false,
+			expectedCluster:       "bar",
+		},
+	}
 
-	err := fakeClusterGetOptions.Complete(fakeArgs)
-	require.NoError(t, err)
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			fakeClusterGetOptions := NewGetClusterOptions(testStreams)
+
+			*fakeClusterGetOptions.configFlags.Namespace = tc.namespace
+
+			err := fakeClusterGetOptions.Complete(tc.args)
+			require.NoError(t, err)
 
-	assert.True(t, fakeClusterGetOptions.AllNamespaces)
-	assert.Equal(t, fakeClusterGetOptions.args, fakeArgs)
+			assert.Equal(t, tc.expectedAllNamespaces, fakeClusterGetOptions.AllNamespaces)
+			assert.Equal(t, tc.expectedCluster, fakeClusterGetOptions.cluster)
+		})
+	}
 }
 
 // Test the Validation() step of the command.
@@ -64,7 +104,7 @@ func TestRayClusterGetValidate(t *testing.T) {
 					KubeConfig: &kubeConfigWithoutCurrentContext,
 				},
 				AllNamespaces: false,
-				args:          []string{"random_arg"},
+				cluster:       "random_arg",
 				ioStreams:     &testStreams,
 			},
 			expectError: "no context is currently set, use \"--context\" or \"kubectl config use-context <context>\" to select a new one",
@@ -98,24 +138,6 @@ func TestRayClusterGetValidate(t *testing.T) {
 				ioStreams: &testStreams,
 			},
 		},
-		{
-			name: "Test validation when more than 1 arg",
-			opts: &GetClusterOptions{
-				// Use fake config to bypass the config flag checks
-				configFlags: &genericclioptions.ConfigFlags{
-					Namespace:        &testNS,
-					Context:          &testContext,
-					KubeConfig:       &kubeConfigWithCurrentContext,
-					BearerToken:      &testBT,
-					Impersonate:      &testImpersonate,
-					ImpersonateGroup: &[]string{"fake-group"},
-				},
-				AllNamespaces: false,
-				args:          []string{"fake", "args"},
-				ioStreams:     &testStreams,
-			},
-			expectError: "too many arguments, either one or no arguments are allowed",
-		},
 		{
 			name: "Successful validation call",
 			opts: &GetClusterOptions{
@@ -129,7 +151,7 @@ func TestRayClusterGetValidate(t *testing.T) {
 					ImpersonateGroup: &[]string{"fake-group"},
 				},
 				AllNamespaces: false,
-				args:          []string{"random_arg"},
+				cluster:       "random_arg",
 				ioStreams:     &testStreams,
 			},
 			expectError: "",