Validate effort level in ParseFlags and add comprehensive tests for valid and invalid effort inputs

Author: pelikhan

cmd/generate/generate.go

@@ -119,6 +119,10 @@ func ParseFlags(cmd *cobra.Command, options *PromptPexOptions) error {
 	flags := cmd.Flags()
 	// Parse effort first so it can set defaults
 	if effort, _ := flags.GetString("effort"); effort != "" {
+		// Validate effort value
+		if effort != EffortLow && effort != EffortMedium && effort != EffortHigh {
+			return fmt.Errorf("invalid effort level '%s': must be one of %s, %s, or %s", effort, EffortLow, EffortMedium, EffortHigh)
+		}
 		options.Effort = effort
 	}
 

cmd/generate/generate_test.go

@@ -75,6 +75,20 @@ func TestParseFlags(t *testing.T) {
 				require.Equal(t, "medium", opts.Effort)
 			},
 		},
+		{
+			name: "valid effort low",
+			args: []string{"--effort", "low"},
+			validate: func(t *testing.T, opts *PromptPexOptions) {
+				require.Equal(t, "low", opts.Effort)
+			},
+		},
+		{
+			name: "valid effort high",
+			args: []string{"--effort", "high"},
+			validate: func(t *testing.T, opts *PromptPexOptions) {
+				require.Equal(t, "high", opts.Effort)
+			},
+		},
 		{
 			name: "groundtruth model flag",
 			args: []string{"--groundtruth-model", "openai/gpt-4o"},
@@ -157,6 +171,63 @@ func TestParseFlags(t *testing.T) {
 	}
 }
 
+func TestParseFlagsInvalidEffort(t *testing.T) {
+	tests := []struct {
+		name        string
+		effort      string
+		expectedErr string
+	}{
+		{
+			name:        "invalid effort value",
+			effort:      "invalid",
+			expectedErr: "invalid effort level 'invalid': must be one of low, medium, or high",
+		},
+		{
+			name:        "empty effort value",
+			effort:      "",
+			expectedErr: "", // Empty should be allowed (no error)
+		},
+		{
+			name:        "case sensitive effort",
+			effort:      "Low",
+			expectedErr: "invalid effort level 'Low': must be one of low, medium, or high",
+		},
+		{
+			name:        "numeric effort",
+			effort:      "1",
+			expectedErr: "invalid effort level '1': must be one of low, medium, or high",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// Create a temporary command to parse flags
+			cmd := NewGenerateCommand(nil)
+			args := []string{}
+			if tt.effort != "" {
+				args = append(args, "--effort", tt.effort)
+			}
+			args = append(args, "dummy.yml") // Add required positional arg
+			cmd.SetArgs(args)
+
+			// Parse flags but don't execute
+			err := cmd.ParseFlags(args[:len(args)-1]) // Exclude positional arg from flag parsing
+			require.NoError(t, err)
+
+			// Parse options from the flags
+			options := GetDefaultOptions()
+			err = ParseFlags(cmd, options)
+
+			if tt.expectedErr == "" {
+				require.NoError(t, err)
+			} else {
+				require.Error(t, err)
+				require.Contains(t, err.Error(), tt.expectedErr)
+			}
+		})
+	}
+}
+
 func TestGenerateCommandExecution(t *testing.T) {
 
 	t.Run("fails with invalid prompt file", func(t *testing.T) {