diff --git a/policy/policy.go b/policy/policy.go index 4472fb90..de72beb8 100644 --- a/policy/policy.go +++ b/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) { diff --git a/policy/risk_factor.go b/policy/risk_factor.go index 275bb06d..0829a591 100644 --- a/policy/risk_factor.go +++ b/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) + } + } } diff --git a/policy/risk_factor_test.go b/policy/risk_factor_test.go index 8e783bd2..58ac65ce 100644 --- a/policy/risk_factor_test.go +++ b/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) +}