🐛 dedupe ScoredRiskFactors.Add (#1234)

Signed-off-by: Dominik Richter <[email protected]>

policy/policy.go

@@ -56,11 +56,10 @@ func (dqi *DataQueryInfo) UnmarshalJSON(data []byte) error {
 
 // WaitUntilDone for a score and an entity
 func WaitUntilDone(resolver PolicyResolver, entity string, scoringMrn string, timeout time.Duration) (bool, error) {
-	var found bool
 	start := time.Now()
 	ctx := context.Background()
 
-	for time.Now().Sub(start) < timeout {
+	for time.Since(start) < timeout {
 		res, err := resolver.GetScore(ctx, &EntityScoreReq{
 			EntityMrn: entity,
 			ScoreMrn:  scoringMrn,
@@ -70,7 +69,6 @@ func WaitUntilDone(resolver PolicyResolver, entity string, scoringMrn string, ti
 		}
 
 		if res != nil && res.Score.ScoreCompletion == 100 && res.Score.DataCompletion == 100 {
-			found = true
 			log.Debug().
 				Str("asset", entity).
 				Str("type", res.Score.TypeLabel()).
@@ -79,13 +77,13 @@ func WaitUntilDone(resolver PolicyResolver, entity string, scoringMrn string, ti
 				Int("data-completion", int(res.Score.DataCompletion)).
 				Int("data-total", int(res.Score.DataTotal)).
 				Msg("waituntildone> got entity score")
-			break
+			return true, nil
 		}
 
 		time.Sleep(50 * time.Millisecond)
 	}
 
-	return found, nil
+	return false, nil
 }
 
 func cannotLookupFilters(ctx context.Context, mrn string) (*explorer.Mquery, error) {

policy/risk_factor.go

@@ -193,5 +193,21 @@ func (s *ScoredRiskFactors) Add(other *ScoredRiskFactors) {
 	if other == nil {
 		return
 	}
-	s.Items = append(s.Items, other.Items...)
+
+	for i := range other.Items {
+		nu := other.Items[i]
+
+		found := false
+		for j := range s.Items {
+			if s.Items[j].Mrn == nu.Mrn {
+				s.Items[j] = nu
+				found = true
+				break
+			}
+		}
+
+		if !found {
+			s.Items = append(s.Items, nu)
+		}
+	}
 }

policy/risk_factor_test.go

@@ -281,3 +281,25 @@ func TestRiskFactor_AdjustRiskScore(t *testing.T) {
 		})
 	}
 }
+
+func TestScoredRiskFactors_Add(t *testing.T) {
+	risks := &ScoredRiskFactors{}
+	risks.Add(&ScoredRiskFactors{
+		Items: []*ScoredRiskFactor{
+			{Mrn: "//mrn1", Risk: -0.2},
+			{Mrn: "//mrn2", Risk: -0.4},
+		},
+	})
+	risks.Add(&ScoredRiskFactors{
+		Items: []*ScoredRiskFactor{
+			{Mrn: "//mrn1", Risk: -0.6},
+			{Mrn: "//mrn3", Risk: -0.9},
+		},
+	})
+
+	assert.Equal(t, []*ScoredRiskFactor{
+		{Mrn: "//mrn1", Risk: -0.6},
+		{Mrn: "//mrn2", Risk: -0.4},
+		{Mrn: "//mrn3", Risk: -0.9},
+	}, risks.Items)
+}