[refactor][kubectl-plugin] make --node-type an enum flag (ray-proje…

…ct#3018) Signed-off-by: David Xia <david@davidxia.com>
davidxia · Feb 18, 2025 · 5cf6bf2 · 5cf6bf2
1 parent e0c7744
commit 5cf6bf2
Show file tree

Hide file tree

Showing 2 changed files with 55 additions and 23 deletions.
diff --git a/kubectl-plugin/pkg/cmd/log/log.go b/kubectl-plugin/pkg/cmd/log/log.go
@@ -32,12 +32,51 @@ import (
 
 const filePathInPod = "/tmp/ray/session_latest/logs/"
 
+type nodeTypeEnum string
+
+const (
+	allNodeType    nodeTypeEnum = "all"
+	headNodeType   nodeTypeEnum = "head"
+	workerNodeType nodeTypeEnum = "worker"
+)
+
+// String is used both by fmt.Print and by Cobra in help text
+func (e *nodeTypeEnum) String() string {
+	return string(*e)
+}
+
+// Set must have pointer receiver so it doesn't change the value of a copy
+func (e *nodeTypeEnum) Set(v string) error {
+	val := strings.ToLower(v)
+
+	switch val {
+	case string(allNodeType), string(headNodeType), string(workerNodeType):
+		*e = nodeTypeEnum(val)
+		return nil
+	default:
+		return fmt.Errorf("must be one of %q, %q, or %q", allNodeType, headNodeType, workerNodeType)
+	}
+}
+
+// Type is only used in help text
+func (e *nodeTypeEnum) Type() string {
+	return "enum"
+}
+
+func nodeTypeCompletion(_ *cobra.Command, _ []string, _ string) ([]string, cobra.ShellCompDirective) {
+	return []string{
+		fmt.Sprintf("%s\tdownload logs from %s Ray nodes", allNodeType, allNodeType),
+		fmt.Sprintf("%s\tdownload logs from %s Ray nodes", headNodeType, headNodeType),
+		fmt.Sprintf("%s\tdownload logs from %s Ray nodes", workerNodeType, workerNodeType),
+	}, cobra.ShellCompDirectiveDefault
+}
+
 type ClusterLogOptions struct {
 	configFlags  *genericclioptions.ConfigFlags
 	ioStreams    *genericclioptions.IOStreams
 	Executor     RemoteExecutor
 	outputDir    string
-	nodeType     string
+	nodeType     nodeTypeEnum
 	ResourceName string
 	ResourceType util.ResourceType
 }
@@ -73,6 +112,7 @@ func NewClusterLogOptions(streams genericclioptions.IOStreams) *ClusterLogOption
 		configFlags: genericclioptions.NewConfigFlags(true),
 		ioStreams:   &streams,
 		Executor:    &DefaultRemoteExecutor{},
+		nodeType:    allNodeType,
 	}
 }
 
@@ -106,7 +146,10 @@ func NewClusterLogCommand(streams genericclioptions.IOStreams) *cobra.Command {
 		},
 	}
 	cmd.Flags().StringVar(&options.outputDir, "out-dir", options.outputDir, "directory to save the logs to")
-	cmd.Flags().StringVar(&options.nodeType, "node-type", options.nodeType, "type of Ray node from which to download log, supports 'worker', 'head', or 'all'")
+	cmd.Flags().Var(&options.nodeType, "node-type", "type of Ray node from which to download logs, supports: 'worker', 'head', or 'all'")
+
+	cobra.CheckErr(cmd.RegisterFlagCompletionFunc("node-type", nodeTypeCompletion))
+
 	options.configFlags.AddFlags(cmd.Flags())
 	return cmd
 }
@@ -139,12 +182,6 @@ func (options *ClusterLogOptions) Complete(cmd *cobra.Command, args []string) er
 		options.ResourceName = typeAndName[1]
 	}
 
-	if options.nodeType == "" {
-		options.nodeType = "all"
-	} else {
-		options.nodeType = strings.ToLower(options.nodeType)
-	}
-
 	return nil
 }
 

diff --git a/kubectl-plugin/pkg/cmd/log/log_test.go b/kubectl-plugin/pkg/cmd/log/log_test.go
@@ -117,18 +117,16 @@ func TestRayClusterLogComplete(t *testing.T) {
 
 	tests := []struct {
 		name                 string
-		nodeType             string
+		nodeType             nodeTypeEnum
 		expectedResourceType util.ResourceType
 		expectedResourceName string
-		expectedNodeType     string
 		args                 []string
 		hasErr               bool
 	}{
 		{
-			name:                 "valid request with RayCluster with empty nodetype input",
+			name:                 "valid request with RayCluster",
 			expectedResourceType: util.RayCluster,
 			expectedResourceName: "test-raycluster",
-			expectedNodeType:     "all",
 			args:                 []string{"test-raycluster"},
 			hasErr:               false,
 		},
@@ -137,23 +135,20 @@ func TestRayClusterLogComplete(t *testing.T) {
 			expectedResourceType: util.RayCluster,
 			expectedResourceName: "test-raycluster",
 			args:                 []string{"rayCluster/test-raycluster"},
-			expectedNodeType:     "all",
 			hasErr:               false,
 		},
 		{
 			name:                 "valid request with RayService",
 			expectedResourceType: util.RayService,
 			expectedResourceName: "test-rayservice",
 			args:                 []string{"rayservice/test-rayservice"},
-			expectedNodeType:     "all",
 			hasErr:               false,
 		},
 		{
 			name:                 "valid request with RayJob",
 			expectedResourceType: util.RayJob,
 			expectedResourceName: "test-rayJob",
 			args:                 []string{"rayJob/test-rayJob"},
-			expectedNodeType:     "all",
 			hasErr:               false,
 		},
 		{
@@ -177,14 +172,14 @@ func TestRayClusterLogComplete(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			testStreams, _, _, _ := genericclioptions.NewTestIOStreams()
 			fakeClusterLogOptions := NewClusterLogOptions(testStreams)
+			fakeClusterLogOptions.nodeType = tc.nodeType
 			err := fakeClusterLogOptions.Complete(cmd, tc.args)
 			if tc.hasErr {
 				require.Error(t, err)
 			} else {
 				require.NoError(t, err)
 				assert.Equal(t, tc.expectedResourceType, fakeClusterLogOptions.ResourceType)
 				assert.Equal(t, tc.expectedResourceName, fakeClusterLogOptions.ResourceName)
-				assert.Equal(t, tc.expectedNodeType, fakeClusterLogOptions.nodeType)
 			}
 		})
 	}
@@ -227,7 +222,7 @@ func TestRayClusterLogValidate(t *testing.T) {
 				},
 				outputDir:    fakeDir,
 				ResourceName: "fake-cluster",
-				nodeType:     "head",
+				nodeType:     headNodeType,
 				ioStreams:    &testStreams,
 			},
 			expectError: "no context is currently set, use \"--context\" or \"kubectl config use-context <context>\" to select a new one",
@@ -241,7 +236,7 @@ func TestRayClusterLogValidate(t *testing.T) {
 				},
 				outputDir:    fakeDir,
 				ResourceName: "fake-cluster",
-				nodeType:     "head",
+				nodeType:     workerNodeType,
 				ioStreams:    &testStreams,
 			},
 		},
@@ -254,7 +249,7 @@ func TestRayClusterLogValidate(t *testing.T) {
 				},
 				outputDir:    fakeDir,
 				ResourceName: "fake-cluster",
-				nodeType:     "head",
+				nodeType:     allNodeType,
 				ioStreams:    &testStreams,
 			},
 		},
@@ -277,7 +272,7 @@ func TestRayClusterLogValidate(t *testing.T) {
 				configFlags:  fakeConfigFlags,
 				outputDir:    fakeDir,
 				ResourceName: "fake-cluster",
-				nodeType:     "head",
+				nodeType:     headNodeType,
 				ioStreams:    &testStreams,
 			},
 			expectError: "",
@@ -289,7 +284,7 @@ func TestRayClusterLogValidate(t *testing.T) {
 				configFlags:  fakeConfigFlags,
 				outputDir:    "",
 				ResourceName: "fake-cluster",
-				nodeType:     "head",
+				nodeType:     headNodeType,
 				ioStreams:    &testStreams,
 			},
 			expectError: "",
@@ -301,7 +296,7 @@ func TestRayClusterLogValidate(t *testing.T) {
 				configFlags:  fakeConfigFlags,
 				outputDir:    "randomPath-here",
 				ResourceName: "fake-cluster",
-				nodeType:     "head",
+				nodeType:     headNodeType,
 				ioStreams:    &testStreams,
 			},
 			expectError: "Directory does not exist. Failed with: stat randomPath-here: no such file or directory",
@@ -313,7 +308,7 @@ func TestRayClusterLogValidate(t *testing.T) {
 				configFlags:  fakeConfigFlags,
 				outputDir:    kubeConfigWithCurrentContext,
 				ResourceName: "fake-cluster",
-				nodeType:     "head",
+				nodeType:     headNodeType,
 				ioStreams:    &testStreams,
 			},
 			expectError: "Path is not a directory. Please input a directory and try again",