Skip to content

Commit

Permalink
[refactor][kubectl-plugin] use cobra's posargs length validator (#2985)
Browse files Browse the repository at this point in the history
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]>
  • Loading branch information
davidxia authored Feb 13, 2025
1 parent 39e8028 commit 3e931a2
Show file tree
Hide file tree
Showing 2 changed files with 69 additions and 50 deletions.
23 changes: 12 additions & 11 deletions kubectl-plugin/pkg/cmd/get/get_cluster.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ import (
type GetClusterOptions struct {
configFlags *genericclioptions.ConfigFlags
ioStreams *genericclioptions.IOStreams
args []string
cluster string
AllNamespaces bool
}

Expand All @@ -44,6 +44,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
Expand All @@ -69,7 +70,10 @@ func (options *GetClusterOptions) Complete(args []string) error {
options.AllNamespaces = true
}

options.args = args
if len(args) >= 1 {
options.cluster = args[0]
}

return nil
}

Expand All @@ -82,9 +86,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
}

Expand All @@ -102,26 +103,26 @@ func getRayClusters(ctx context.Context, options *GetClusterOptions, k8sClient c
var err error

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 nil, fmt.Errorf("unable to retrieve raycluster for all namespaces: %w", err)
return nil, 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 nil, fmt.Errorf("unable to retrieve raycluster for namespace %s: %w", *options.configFlags.Namespace, err)
return nil, fmt.Errorf("unable to retrieve Ray clusters for namespace %s: %w", *options.configFlags.Namespace, err)
}
}

if len(options.args) == 1 && len(rayclusterList.Items) == 0 {
errMsg := fmt.Sprintf("Ray cluster %s not found", options.args[0])
if options.cluster != "" && len(rayclusterList.Items) == 0 {
errMsg := fmt.Sprintf("Ray cluster %s not found", options.cluster)
if options.AllNamespaces {
errMsg += " in any namespace"
} else {
Expand Down
96 changes: 57 additions & 39 deletions kubectl-plugin/pkg/cmd/get/get_cluster_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -24,22 +24,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",
},
}

for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
fakeClusterGetOptions := NewGetClusterOptions(testStreams)

err := fakeClusterGetOptions.Complete(fakeArgs)
require.NoError(t, err)
*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.
Expand Down Expand Up @@ -67,7 +107,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",
Expand Down Expand Up @@ -101,24 +141,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{
Expand All @@ -132,7 +154,7 @@ func TestRayClusterGetValidate(t *testing.T) {
ImpersonateGroup: &[]string{"fake-group"},
},
AllNamespaces: false,
args: []string{"random_arg"},
cluster: "random_arg",
ioStreams: &testStreams,
},
expectError: "",
Expand Down Expand Up @@ -236,56 +258,52 @@ func TestGetRayClusters(t *testing.T) {
name string
expectedError string
expectedOutput string
args []string
cluster string
rayClusters []runtime.Object
}{
{
name: "should not error if no cluster name is provided, searching all namespaces, and no clusters are found",
args: []string{},
namespace: nil,
rayClusters: []runtime.Object{},
},
{
name: "should error if a cluster name is provided, searching all namespaces, and no clusters are found",
args: []string{"my-cluster"},
cluster: "my-cluster",
namespace: nil,
rayClusters: []runtime.Object{},
expectedError: "Ray cluster my-cluster not found",
},
{
name: "should not error if no cluster name is provided, searching one namespace, and no clusters are found",
args: []string{},
namespace: &namespace,
rayClusters: []runtime.Object{},
},
{
name: "should error if a cluster name is provided, searching one namespace, and no clusters are found",
args: []string{"my-cluster"},
cluster: "my-cluster",
namespace: &namespace,
rayClusters: []runtime.Object{},
expectedError: fmt.Sprintf("Ray cluster my-cluster not found in namespace %s", namespace),
},
{
name: "should not error if no cluster name is provided, searching all namespaces, and clusters are found",
args: []string{},
namespace: nil,
rayClusters: []runtime.Object{rayCluster},
},
{
name: "should not error if a cluster name is provided, searching all namespaces, and clusters are found",
args: []string{"my-cluster"},
cluster: "my-cluster",
namespace: nil,
rayClusters: []runtime.Object{rayCluster},
},
{
name: "should not error if no cluster name is provided, searching one namespace, and clusters are found",
args: []string{},
namespace: &namespace,
rayClusters: []runtime.Object{rayCluster},
},
{
name: "should not error if a cluster name is provided, searching one namespace, and clusters are found",
args: []string{"my-cluster"},
cluster: "my-cluster",
namespace: &namespace,
rayClusters: []runtime.Object{rayCluster},
},
Expand All @@ -296,7 +314,7 @@ func TestGetRayClusters(t *testing.T) {
fakeClusterGetOptions := GetClusterOptions{
configFlags: genericclioptions.NewConfigFlags(true),
ioStreams: &testStreams,
args: tc.args,
cluster: tc.cluster,
}
if tc.namespace != nil {
*fakeClusterGetOptions.configFlags.Namespace = *tc.namespace
Expand Down

0 comments on commit 3e931a2

Please sign in to comment.