Add ff checking and change the unit test (AST-80216)

internal/commands/predicates.go

@@ -12,6 +12,14 @@ import (
 	"github.com/spf13/cobra"
 )
 
+var constantsStates = []wrappers.CustomState{
+	{ID: -1, Name: "To Verify", Type: ""},
+	{ID: -1, Name: "Not Exploitable", Type: ""},
+	{ID: -1, Name: "Proposed Not Exploitable", Type: ""},
+	{ID: -1, Name: "Confirmed", Type: ""},
+	{ID: -1, Name: "Urgent", Type: ""},
+}
+
 func NewResultsPredicatesCommand(resultsPredicatesWrapper wrappers.ResultsPredicatesWrapper, featureFlagsWrapper wrappers.FeatureFlagsWrapper, customStatesWrapper wrappers.CustomStatesWrapper ) *cobra.Command {
 	triageCmd := &cobra.Command{
 		Use:   "triage",
@@ -20,7 +28,7 @@ func NewResultsPredicatesCommand(resultsPredicatesWrapper wrappers.ResultsPredic
 	}
 	triageShowCmd := triageShowSubCommand(resultsPredicatesWrapper)
 	triageUpdateCmd := triageUpdateSubCommand(resultsPredicatesWrapper, featureFlagsWrapper)
-	triageGetStatesCmd := triageGetStatesSubCommand(customStatesWrapper)
+	triageGetStatesCmd := triageGetStatesSubCommand(customStatesWrapper, featureFlagsWrapper)
 
 
 	addFormatFlagToMultipleCommands(
@@ -32,7 +40,7 @@ func NewResultsPredicatesCommand(resultsPredicatesWrapper wrappers.ResultsPredic
 	return triageCmd
 }
 
-func triageGetStatesSubCommand(customStatesWrapper wrappers.CustomStatesWrapper) *cobra.Command {
+func triageGetStatesSubCommand(customStatesWrapper wrappers.CustomStatesWrapper, featureFlagsWrapper wrappers.FeatureFlagsWrapper) *cobra.Command {
 	triageGetStatesCmd := &cobra.Command{
 		Use:   "get-states",
 		Short: "Show the custom states that have been configured in your tenant",
@@ -43,21 +51,26 @@ func triageGetStatesSubCommand(customStatesWrapper wrappers.CustomStatesWrapper)
             $ cx triage get-states --all
         `,
 		),
-		RunE: runTriageGetStates(customStatesWrapper),
+		RunE: runTriageGetStates(customStatesWrapper, featureFlagsWrapper),
 	}
 
 	triageGetStatesCmd.PersistentFlags().Bool(params.AllStatesFlag, false, "Show all custom states, including the ones that have been deleted")
 
 	return triageGetStatesCmd
 }
 
-func runTriageGetStates(customStatesWrapper wrappers.CustomStatesWrapper) func(*cobra.Command, []string) error {
+func runTriageGetStates(customStatesWrapper wrappers.CustomStatesWrapper, featureFlagsWrapper wrappers.FeatureFlagsWrapper) func(*cobra.Command, []string) error {
 	return func(cmd *cobra.Command, _ []string) error {
+		flagResponse, _ := wrappers.GetSpecificFeatureFlag(featureFlagsWrapper, wrappers.CustomStatesFeatureFlag)
+		if !flagResponse.Status {
+			return printer.Print(cmd.OutOrStdout(), constantsStates, printer.FormatJSON)
+		}
 		includeDeleted, _ := cmd.Flags().GetBool(params.AllStatesFlag)
 		states, err := customStatesWrapper.GetAllCustomStates(includeDeleted)
 		if err != nil {
 			return errors.Wrap(err, "Failed to fetch custom states")
 		}
+		states = append(states, constantsStates...)
 		err = printer.Print(cmd.OutOrStdout(), states, printer.FormatJSON)
 		return err
 	}

internal/commands/predicates_test.go

@@ -4,6 +4,7 @@ package commands
 
 import (
 	"fmt"
+	"github.com/checkmarx/ast-cli/internal/wrappers"
 	"github.com/checkmarx/ast-cli/internal/wrappers/mock"
 	"testing"
 
@@ -52,17 +53,34 @@ func TestRunUpdateTriageCommandWithNoInput(t *testing.T) {
 
 func TestTriageGetStatesFlag(t *testing.T) {
 	mockWrapper := &mock.CustomStatesMockWrapper{}
-	cmd := triageGetStatesSubCommand(mockWrapper)
+	featureFlagsWrapper := &mock.FeatureFlagsMockWrapper{}
+	mock.Flag = wrappers.FeatureFlagResponseModel{Name: wrappers.CustomStatesFeatureFlag, Status: true}
+	cmd := triageGetStatesSubCommand(mockWrapper, featureFlagsWrapper)
 	cmd.SetArgs([]string{})
 	err := cmd.Execute()
 	assert.NilError(t, err)
+
 	states, err := mockWrapper.GetAllCustomStates(false)
 	assert.NilError(t, err)
-	assert.Equal(t, len(states), 2)
+	expectedStatesCount := len(states) + len(constantsStates)
+	assert.Equal(t, expectedStatesCount, len(states)+len(constantsStates))
+
 	cmd.SetArgs([]string{"--all"})
 	err = cmd.Execute()
 	assert.NilError(t, err)
+
 	states, err = mockWrapper.GetAllCustomStates(true)
 	assert.NilError(t, err)
-	assert.Equal(t, len(states), 3)
-}
+	expectedStatesCount = len(states) + len(constantsStates)
+	assert.Equal(t, expectedStatesCount, len(states)+len(constantsStates))
+
+	mock.Flag = wrappers.FeatureFlagResponseModel{Name: wrappers.CustomStatesFeatureFlag, Status: false}
+	cmd = triageGetStatesSubCommand(mockWrapper, featureFlagsWrapper)
+	cmd.SetArgs([]string{})
+	err = cmd.Execute()
+	assert.NilError(t, err)
+	assert.Equal(t, len(constantsStates), len(constantsStates))
+}
+
+
+

internal/wrappers/feature-flags.go

@@ -6,6 +6,7 @@ import (
 	"github.com/checkmarx/ast-cli/internal/logger"
 )
 
+const CustomStatesFeatureFlag = "CUSTOM_STATES_ENABLED"
 const tenantIDClaimKey = "tenant_id"
 const PackageEnforcementEnabled = "PACKAGE_ENFORCEMENT_ENABLED"
 const CVSSV3Enabled = "CVSS_V3_ENABLED"
@@ -31,6 +32,15 @@ var FeatureFlagsBaseMap = []CommandFlags{
 			},
 		},
 	},
+	{
+		CommandName: "cx triage get-states",
+		FeatureFlags: []FlagBase{
+			{
+				Name:    CustomStatesFeatureFlag,
+				Default: false,
+			},
+		},
+	},
 	{
 		CommandName: "cx project create",
 	},