diff --git a/internal/commands/predicates.go b/internal/commands/predicates.go index f0501291a..c292f8f3e 100644 --- a/internal/commands/predicates.go +++ b/internal/commands/predicates.go @@ -59,16 +59,36 @@ func triageGetStatesSubCommand(customStatesWrapper wrappers.CustomStatesWrapper, } triageGetStatesCmd.PersistentFlags().Bool(params.AllStatesFlag, false, "Show all custom states, including the ones that have been deleted") - + triageGetStatesCmd.PersistentFlags().String(params.ScanTypeFlag, "sast", "Scan Type") return triageGetStatesCmd } func runTriageGetStates(customStatesWrapper wrappers.CustomStatesWrapper, featureFlagsWrapper wrappers.FeatureFlagsWrapper) func(*cobra.Command, []string) error { return func(cmd *cobra.Command, _ []string) error { + scanType, _ := cmd.Flags().GetString(params.ScanTypeFlag) + scanType = strings.TrimSpace(strings.ToLower(scanType)) flagResponse, _ := wrappers.GetSpecificFeatureFlag(featureFlagsWrapper, wrappers.CustomStatesFeatureFlag) if !flagResponse.Status { return printer.Print(cmd.OutOrStdout(), systemStates, printer.FormatJSON) } + + switch scanType { + case params.KicsType: + kicsCustomStates, _ := wrappers.GetSpecificFeatureFlag(featureFlagsWrapper, "KICS_CUSTOM_STATES_ENABLED") + if !kicsCustomStates.Status { + return printer.Print(cmd.OutOrStdout(), systemStates, printer.FormatJSON) + } + case params.ScaType: + scaCustomStates, _ := wrappers.GetSpecificFeatureFlag(featureFlagsWrapper, "SCA_CUSTOM_STATES") + if !scaCustomStates.Status { + return printer.Print(cmd.OutOrStdout(), systemStates, printer.FormatJSON) + } + default: + sastCustomStates, _ := wrappers.GetSpecificFeatureFlag(featureFlagsWrapper, "SAST_CUSTOM_STATES_ENABLED") + if !sastCustomStates.Status { + return printer.Print(cmd.OutOrStdout(), systemStates, printer.FormatJSON) + } + } includeDeleted, _ := cmd.Flags().GetBool(params.AllStatesFlag) states, err := customStatesWrapper.GetAllCustomStates(includeDeleted) if err != nil { diff --git a/internal/commands/predicates_test.go b/internal/commands/predicates_test.go index c7d49461c..913d39e1f 100644 --- a/internal/commands/predicates_test.go +++ b/internal/commands/predicates_test.go @@ -213,6 +213,187 @@ func TestRunTriageUpdateWithSystemState(t *testing.T) { assert.NilError(t, err) } +func TestRunTriageGetStates_ScanTypeSwitch(t *testing.T) { + tests := []struct { + name string + scanType string + customStatesFeatureFlag wrappers.FeatureFlagResponseModel + scanTypeFeatureFlag wrappers.FeatureFlagResponseModel + expectedError bool + description string + }{ + { + name: "KICS scan type with KICS custom states enabled", + scanType: "kics", + customStatesFeatureFlag: wrappers.FeatureFlagResponseModel{Name: wrappers.CustomStatesFeatureFlag, Status: true}, + scanTypeFeatureFlag: wrappers.FeatureFlagResponseModel{Name: "KICS_CUSTOM_STATES_ENABLED", Status: true}, + expectedError: false, + description: "Should successfully get custom states for KICS when feature flag is enabled", + }, + { + name: "KICS scan type with KICS custom states disabled", + scanType: "kics", + customStatesFeatureFlag: wrappers.FeatureFlagResponseModel{Name: wrappers.CustomStatesFeatureFlag, Status: true}, + scanTypeFeatureFlag: wrappers.FeatureFlagResponseModel{Name: "KICS_CUSTOM_STATES_ENABLED", Status: false}, + expectedError: false, + description: "Should return system states for KICS when feature flag is disabled", + }, + { + name: "SCA scan type with SCA custom states enabled", + scanType: "sca", + customStatesFeatureFlag: wrappers.FeatureFlagResponseModel{Name: wrappers.CustomStatesFeatureFlag, Status: true}, + scanTypeFeatureFlag: wrappers.FeatureFlagResponseModel{Name: "SCA_CUSTOM_STATES", Status: true}, + expectedError: false, + description: "Should successfully get custom states for SCA when feature flag is enabled", + }, + { + name: "SCA scan type with SCA custom states disabled", + scanType: "sca", + customStatesFeatureFlag: wrappers.FeatureFlagResponseModel{Name: wrappers.CustomStatesFeatureFlag, Status: true}, + scanTypeFeatureFlag: wrappers.FeatureFlagResponseModel{Name: "SCA_CUSTOM_STATES", Status: false}, + expectedError: false, + description: "Should return system states for SCA when feature flag is disabled", + }, + { + name: "SAST scan type with SAST custom states enabled", + scanType: "sast", + customStatesFeatureFlag: wrappers.FeatureFlagResponseModel{Name: wrappers.CustomStatesFeatureFlag, Status: true}, + scanTypeFeatureFlag: wrappers.FeatureFlagResponseModel{Name: "SAST_CUSTOM_STATES_ENABLED", Status: true}, + expectedError: false, + description: "Should successfully get custom states for SAST when feature flag is enabled", + }, + { + name: "SAST scan type with SAST custom states disabled", + scanType: "sast", + customStatesFeatureFlag: wrappers.FeatureFlagResponseModel{Name: wrappers.CustomStatesFeatureFlag, Status: true}, + scanTypeFeatureFlag: wrappers.FeatureFlagResponseModel{Name: "SAST_CUSTOM_STATES_ENABLED", Status: false}, + expectedError: false, + description: "Should return system states for SAST when feature flag is disabled", + }, + { + name: "Unknown scan type defaults to SAST behavior", + scanType: "unknown", + customStatesFeatureFlag: wrappers.FeatureFlagResponseModel{Name: wrappers.CustomStatesFeatureFlag, Status: true}, + scanTypeFeatureFlag: wrappers.FeatureFlagResponseModel{Name: "SAST_CUSTOM_STATES_ENABLED", Status: true}, + expectedError: false, + description: "Should default to SAST logic for unknown scan types", + }, + { + name: "KICS with uppercase scan type", + scanType: "KICS", + customStatesFeatureFlag: wrappers.FeatureFlagResponseModel{Name: wrappers.CustomStatesFeatureFlag, Status: true}, + scanTypeFeatureFlag: wrappers.FeatureFlagResponseModel{Name: "KICS_CUSTOM_STATES_ENABLED", Status: true}, + expectedError: false, + description: "Should handle uppercase scan type correctly", + }, + { + name: "SCA with mixed case scan type", + scanType: "ScA", + customStatesFeatureFlag: wrappers.FeatureFlagResponseModel{Name: wrappers.CustomStatesFeatureFlag, Status: true}, + scanTypeFeatureFlag: wrappers.FeatureFlagResponseModel{Name: "SCA_CUSTOM_STATES", Status: true}, + expectedError: false, + description: "Should handle mixed case scan type correctly", + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + clearFlags() + mockCustomStatesWrapper := &mock.CustomStatesMockWrapper{} + mockFeatureFlagsWrapper := &mock.FeatureFlagsMockWrapper{} + + // First, set the custom states feature flag + mock.Flag = tt.customStatesFeatureFlag + cmd := triageGetStatesSubCommand(mockCustomStatesWrapper, mockFeatureFlagsWrapper) + + // Then set the scan-type specific feature flag for the actual execution + mock.Flag = tt.scanTypeFeatureFlag + cmd.SetArgs([]string{"--scan-type", tt.scanType}) + + err := cmd.Execute() + + if tt.expectedError { + assert.Assert(t, err != nil, tt.description) + } else { + assert.NilError(t, err, tt.description) + } + }) + } +} + +func TestRunTriageGetStates_OverallFeatureFlagDisabled(t *testing.T) { + clearFlags() + mockCustomStatesWrapper := &mock.CustomStatesMockWrapper{} + mockFeatureFlagsWrapper := &mock.FeatureFlagsMockWrapper{} + + // Set the overall custom states feature flag to disabled + mock.Flag = wrappers.FeatureFlagResponseModel{Name: wrappers.CustomStatesFeatureFlag, Status: false} + + cmd := triageGetStatesSubCommand(mockCustomStatesWrapper, mockFeatureFlagsWrapper) + cmd.SetArgs([]string{"--scan-type", "sast"}) + + err := cmd.Execute() + assert.NilError(t, err, "Should return system states when overall custom states feature flag is disabled") +} + +func TestRunTriageGetStates_WithAllFlag(t *testing.T) { + clearFlags() + mockCustomStatesWrapper := &mock.CustomStatesMockWrapper{} + mockFeatureFlagsWrapper := &mock.FeatureFlagsMockWrapper{} + + tests := []struct { + name string + scanType string + allFlag bool + }{ + { + name: "KICS with all flag true", + scanType: "kics", + allFlag: true, + }, + { + name: "SCA with all flag false", + scanType: "sca", + allFlag: false, + }, + { + name: "SAST with all flag true", + scanType: "sast", + allFlag: true, + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + clearFlags() + mock.Flag = wrappers.FeatureFlagResponseModel{Name: wrappers.CustomStatesFeatureFlag, Status: true} + + cmd := triageGetStatesSubCommand(mockCustomStatesWrapper, mockFeatureFlagsWrapper) + + // Set the scan-type specific feature flag to enabled + switch tt.scanType { + case "kics": + mock.Flag = wrappers.FeatureFlagResponseModel{Name: "KICS_CUSTOM_STATES_ENABLED", Status: true} + case "sca": + mock.Flag = wrappers.FeatureFlagResponseModel{Name: "SCA_CUSTOM_STATES", Status: true} + default: + mock.Flag = wrappers.FeatureFlagResponseModel{Name: "SAST_CUSTOM_STATES_ENABLED", Status: true} + } + + args := []string{"--scan-type", tt.scanType} + if tt.allFlag { + args = append(args, "--all") + } + cmd.SetArgs(args) + + err := cmd.Execute() + assert.NilError(t, err) + }) + } +} + func TestDetermineSystemOrCustomState(t *testing.T) { tests := []struct { name string