Skip to content

Commit 3f13e2b

Browse files
tiagobcxtiagobcx
committed
System state or custom state id, and add tests.
1 parent 2842833 commit 3f13e2b

File tree

3 files changed

+209
-12
lines changed

3 files changed

+209
-12
lines changed

‎internal/commands/predicates.go‎

Lines changed: 34 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -119,7 +119,10 @@ func triageUpdateSubCommand(resultsPredicatesWrapper wrappers.ResultsPredicatesW
119119
markFlagAsRequired(triageUpdateCmd, params.SimilarityIDFlag)
120120
markFlagAsRequired(triageUpdateCmd, params.SeverityFlag)
121121
markFlagAsRequired(triageUpdateCmd, params.ProjectIDFlag)
122-
markFlagAsRequired(triageUpdateCmd, params.StateFlag)
122+
flagResponse, _ := wrappers.GetSpecificFeatureFlag(featureFlagsWrapper, wrappers.SastCustomStateEnabled)
123+
if !flagResponse.Status {
124+
markFlagAsRequired(triageUpdateCmd, params.StateFlag)
125+
}
123126
markFlagAsRequired(triageUpdateCmd, params.ScanTypeFlag)
124127

125128
return triageUpdateCmd
@@ -184,16 +187,11 @@ func runTriageUpdate(resultsPredicatesWrapper wrappers.ResultsPredicatesWrapper,
184187
if !criticalEnabled && strings.EqualFold(severity, "critical") {
185188
return errors.Errorf("%s", "Critical severity is not available for your tenant.This severity status will be enabled shortly")
186189
}
187-
if isCustomState(state) {
188-
if customStateID == "" {
189-
var err error
190-
customStateID, err = getCustomStateID(customStatesWrapper, state)
191-
if err != nil {
192-
return errors.Wrapf(err, "Failed to get custom state ID for state: %s", state)
193-
}
194-
}
195-
} else {
196-
customStateID = ""
190+
191+
var err error
192+
state, customStateID, err = determineSystemOrCustomState(customStatesWrapper, featureFlagsWrapper, state, customStateID)
193+
if err != nil {
194+
return err
197195
}
198196

199197
predicate := &wrappers.PredicateRequest{
@@ -205,15 +203,39 @@ func runTriageUpdate(resultsPredicatesWrapper wrappers.ResultsPredicatesWrapper,
205203
Comment: comment,
206204
}
207205

208-
_, err := resultsPredicatesWrapper.PredicateSeverityAndState(predicate, scanType)
206+
_, err = resultsPredicatesWrapper.PredicateSeverityAndState(predicate, scanType)
209207
if err != nil {
210208
return errors.Wrapf(err, "%s", "Failed updating the predicate")
211209
}
212210

213211
return nil
214212
}
215213
}
214+
func determineSystemOrCustomState(customStatesWrapper wrappers.CustomStatesWrapper, featureFlagsWrapper wrappers.FeatureFlagsWrapper, state, customStateID string) (string, string, error) {
215+
if isCustomState(state) {
216+
flagResponse, _ := wrappers.GetSpecificFeatureFlag(featureFlagsWrapper, wrappers.SastCustomStateEnabled)
217+
if !flagResponse.Status {
218+
return "", "", errors.Errorf("%s", "Custom state is not available for your tenant.")
219+
}
220+
221+
if customStateID == "" {
222+
if state == "" {
223+
return "", "", errors.Errorf("state-id is required when state is not provided")
224+
}
225+
var err error
226+
customStateID, err = getCustomStateID(customStatesWrapper, state)
227+
if err != nil {
228+
return "", "", errors.Wrapf(err, "Failed to get custom state ID for state: %s", state)
229+
}
230+
}
231+
return "", customStateID, nil
232+
}
233+
return state, "", nil
234+
}
216235
func isCustomState(state string) bool {
236+
if state == "" {
237+
return true
238+
}
217239
systemStates := []string{"TO_VERIFY", "NOT_EXPLOITABLE", "PROPOSED_NOT_EXPLOITABLE", "CONFIRMED", "URGENT"}
218240
for _, customState := range systemStates {
219241
if strings.EqualFold(state, customState) {

‎internal/commands/predicates_test.go‎

Lines changed: 170 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -134,3 +134,173 @@ func TestIsCustomState(t *testing.T) {
134134
})
135135
}
136136
}
137+
func TestRunTriageUpdateWithNotFoundCustomState(t *testing.T) {
138+
mockResultsPredicatesWrapper := &mock.ResultsPredicatesMockWrapper{}
139+
mockFeatureFlagsWrapper := &mock.FeatureFlagsMockWrapper{}
140+
mockCustomStatesWrapper := &mock.CustomStatesMockWrapper{}
141+
mock.Flag = wrappers.FeatureFlagResponseModel{
142+
Name: "test-flag",
143+
Status: true,
144+
}
145+
cmd := triageUpdateSubCommand(mockResultsPredicatesWrapper, mockFeatureFlagsWrapper, mockCustomStatesWrapper)
146+
cmd.SetArgs([]string{
147+
"--similarity-id", "MOCK",
148+
"--project-id", "MOCK",
149+
"--severity", "low",
150+
"--state", "CUSTOM_STATE_1",
151+
"--scan-type", "sast",
152+
})
153+
154+
err := cmd.Execute()
155+
assert.ErrorContains(t, err, "Failed to get custom state ID for state: CUSTOM_STATE_1: No matching state found for CUSTOM_STATE_1")
156+
}
157+
158+
func TestRunTriageUpdateWithCustomState(t *testing.T) {
159+
mockResultsPredicatesWrapper := &mock.ResultsPredicatesMockWrapper{}
160+
mockFeatureFlagsWrapper := &mock.FeatureFlagsMockWrapper{}
161+
mockCustomStatesWrapper := &mock.CustomStatesMockWrapper{}
162+
mock.Flag = wrappers.FeatureFlagResponseModel{
163+
Name: "test-flag",
164+
Status: true,
165+
}
166+
cmd := triageUpdateSubCommand(mockResultsPredicatesWrapper, mockFeatureFlagsWrapper, mockCustomStatesWrapper)
167+
cmd.SetArgs([]string{
168+
"--similarity-id", "MOCK",
169+
"--project-id", "MOCK",
170+
"--severity", "low",
171+
"--state", "demo2",
172+
"--scan-type", "sast",
173+
})
174+
175+
err := cmd.Execute()
176+
assert.NilError(t, err)
177+
}
178+
179+
func TestRunTriageUpdateWithSystemState(t *testing.T) {
180+
mockResultsPredicatesWrapper := &mock.ResultsPredicatesMockWrapper{}
181+
mockFeatureFlagsWrapper := &mock.FeatureFlagsMockWrapper{}
182+
mockCustomStatesWrapper := &mock.CustomStatesMockWrapper{}
183+
184+
cmd := triageUpdateSubCommand(mockResultsPredicatesWrapper, mockFeatureFlagsWrapper, mockCustomStatesWrapper)
185+
cmd.SetArgs([]string{
186+
"--similarity-id", "MOCK",
187+
"--project-id", "MOCK",
188+
"--severity", "low",
189+
"--state", "TO_VERIFY",
190+
"--scan-type", "sast",
191+
})
192+
193+
err := cmd.Execute()
194+
assert.NilError(t, err)
195+
}
196+
197+
func TestDetermineSystemOrCustomState(t *testing.T) {
198+
tests := []struct {
199+
name string
200+
state string
201+
customStateID string
202+
mockWrapper wrappers.CustomStatesWrapper
203+
mockFeatureFlags wrappers.FeatureFlagsWrapper
204+
flag wrappers.FeatureFlagResponseModel
205+
expectedState string
206+
expectedCustomState string
207+
expectedError string
208+
}{
209+
{
210+
name: "Custom state with valid state name",
211+
state: "demo2",
212+
customStateID: "",
213+
mockWrapper: &mock.CustomStatesMockWrapper{},
214+
mockFeatureFlags: &mock.FeatureFlagsMockWrapper{},
215+
flag: wrappers.FeatureFlagResponseModel{Name: "test-flag", Status: true},
216+
expectedState: "",
217+
expectedCustomState: "2",
218+
expectedError: "",
219+
},
220+
{
221+
name: "Custom state with valid state ID",
222+
state: "",
223+
customStateID: "2",
224+
mockWrapper: &mock.CustomStatesMockWrapper{},
225+
flag: wrappers.FeatureFlagResponseModel{Name: "test-flag", Status: true},
226+
expectedState: "",
227+
expectedCustomState: "2",
228+
expectedError: "",
229+
},
230+
{
231+
name: "System state",
232+
state: "TO_VERIFY",
233+
customStateID: "",
234+
mockWrapper: &mock.CustomStatesMockWrapper{},
235+
flag: wrappers.FeatureFlagResponseModel{Name: "test-flag", Status: true},
236+
expectedState: "TO_VERIFY",
237+
expectedCustomState: "",
238+
expectedError: "",
239+
},
240+
{
241+
name: "State ID required when state is not provided",
242+
state: "",
243+
customStateID: "",
244+
mockWrapper: &mock.CustomStatesMockWrapper{},
245+
flag: wrappers.FeatureFlagResponseModel{Name: "test-flag", Status: true},
246+
expectedState: "",
247+
expectedCustomState: "",
248+
expectedError: "state-id is required when state is not provided",
249+
},
250+
{
251+
name: "Failed to get custom state ID",
252+
state: "INVALID_STATE",
253+
customStateID: "",
254+
mockWrapper: &mock.CustomStatesMockWrapperWithError{},
255+
flag: wrappers.FeatureFlagResponseModel{Name: "test-flag", Status: true},
256+
expectedState: "",
257+
expectedCustomState: "",
258+
expectedError: "Failed to get custom state ID for state: INVALID_STATE",
259+
},
260+
{
261+
name: "Both state and state ID provided - valid custom state",
262+
state: "demo2",
263+
customStateID: "2",
264+
mockWrapper: &mock.CustomStatesMockWrapper{},
265+
flag: wrappers.FeatureFlagResponseModel{Name: "test-flag", Status: true},
266+
expectedState: "",
267+
expectedCustomState: "2",
268+
expectedError: "",
269+
},
270+
{
271+
name: "Both state and state ID provided - valid system state",
272+
state: "TO_VERIFY",
273+
customStateID: "2",
274+
mockWrapper: &mock.CustomStatesMockWrapper{},
275+
flag: wrappers.FeatureFlagResponseModel{Name: "test-flag", Status: true},
276+
expectedState: "TO_VERIFY",
277+
expectedCustomState: "",
278+
expectedError: "",
279+
},
280+
{
281+
name: "Both state and state ID provided - invalid state",
282+
state: "INVALID_STATE",
283+
customStateID: "2",
284+
mockWrapper: &mock.CustomStatesMockWrapperWithError{},
285+
flag: wrappers.FeatureFlagResponseModel{Name: "test-flag", Status: true},
286+
expectedState: "",
287+
expectedCustomState: "2",
288+
expectedError: "",
289+
},
290+
}
291+
292+
for _, tt := range tests {
293+
tt := tt
294+
t.Run(tt.name, func(t *testing.T) {
295+
mock.Flag = tt.flag
296+
state, customStateID, err := determineSystemOrCustomState(tt.mockWrapper, tt.mockFeatureFlags, tt.state, tt.customStateID)
297+
if tt.expectedError != "" {
298+
assert.ErrorContains(t, err, tt.expectedError)
299+
} else {
300+
assert.NilError(t, err)
301+
}
302+
assert.Equal(t, state, tt.expectedState)
303+
assert.Equal(t, customStateID, tt.expectedCustomState)
304+
})
305+
}
306+
}

‎internal/wrappers/feature-flags.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ const tenantIDClaimKey = "tenant_id"
1010
const PackageEnforcementEnabled = "PACKAGE_ENFORCEMENT_ENABLED"
1111
const CVSSV3Enabled = "CVSS_V3_ENABLED"
1212
const MinioEnabled = "MINIO_ENABLED"
13+
const SastCustomStateEnabled = "SAST_CUSTOM_STATES_ENABLED"
1314
const ContainerEngineCLIEnabled = "CONTAINER_ENGINE_CLI_ENABLED"
1415
const SCSEngineCLIEnabled = "NEW_2MS_SCORECARD_RESULTS_CLI_ENABLED"
1516
const NewScanReportEnabled = "NEW_SAST_SCAN_REPORT_ENABLED"
@@ -59,6 +60,10 @@ var FeatureFlagsBaseMap = []CommandFlags{
5960
Name: CVSSV3Enabled,
6061
Default: false,
6162
},
63+
{
64+
Name: SastCustomStateEnabled,
65+
Default: false,
66+
},
6267
},
6368
},
6469
}

0 commit comments

Comments
 (0)