From 4a0a69057b54f84f2d6b8fd2e84c9ab372a7843a Mon Sep 17 00:00:00 2001
From: Jay Mundrawala <jay@mondoo.com>
Date: Thu, 14 Dec 2023 09:54:56 -0600
Subject: [PATCH] =?UTF-8?q?=F0=9F=A7=B9=20Add=20batch=20stat=20APIs=20(#99?=
 =?UTF-8?q?7)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This allows me to add aggregated scores more efficiently to those
structs.
---
 policy/report.go      | 45 +++++++++++++++++++++++++++----------------
 policy/report_test.go | 17 ++++++++++++++++
 2 files changed, 45 insertions(+), 17 deletions(-)

diff --git a/policy/report.go b/policy/report.go
index 6d85c45c..d9c850d0 100644
--- a/policy/report.go
+++ b/policy/report.go
@@ -40,29 +40,32 @@ func (r *Report) ComputeStats(resolved *ResolvedPolicy) {
 
 	r.Stats = &res
 }
-
 func (s *Stats) Add(score *Score) {
-	s.Total++
+	s.AddCount(score, 1)
+}
+
+func (s *Stats) AddCount(score *Score, count uint32) {
+	s.Total += count
 	switch score.Type {
 	case ScoreType_Unknown:
-		s.Unknown++
+		s.Unknown += count
 	case ScoreType_Result:
 		if score.Value < 100 {
-			s.Failed.Add(score)
+			s.Failed.AddCount(score, count)
 
 			if score.Value < s.Worst {
 				s.Worst = score.Value
 			}
 
 		} else {
-			s.Passed.Add(score)
+			s.Passed.AddCount(score, count)
 		}
 	case ScoreType_Error:
-		s.Errors.Add(score)
+		s.Errors.AddCount(score, count)
 	case ScoreType_Skip:
-		s.Skipped++
+		s.Skipped += count
 	case ScoreType_Unscored:
-		s.Unknown++
+		s.Unknown += count
 	default:
 		log.Warn().Uint32("type", score.Type).Str("id", score.QrId).Msg("ran into unknown score type")
 	}
@@ -70,7 +73,11 @@ func (s *Stats) Add(score *Score) {
 
 // this function also handles nil scores and updates the score distribution accordingly
 func (sd *ScoreDistribution) Add(score *Score) {
-	sd.AddRating(score.Rating())
+	sd.AddCount(score, 1)
+}
+
+func (sd *ScoreDistribution) AddCount(score *Score, count uint32) {
+	sd.AddRatingCount(score.Rating(), count)
 }
 
 // this function also handles nil scores and updates the score distribution accordingly
@@ -79,22 +86,26 @@ func (sd *ScoreDistribution) Remove(score *Score) {
 }
 
 func (sd *ScoreDistribution) AddRating(scoreRating ScoreRating) {
-	sd.Total++
+	sd.AddRatingCount(scoreRating, 1)
+}
+
+func (sd *ScoreDistribution) AddRatingCount(scoreRating ScoreRating, count uint32) {
+	sd.Total += count
 	switch scoreRating {
 	case ScoreRating_aPlus, ScoreRating_a, ScoreRating_aMinus:
-		sd.A++
+		sd.A += count
 	case ScoreRating_bPlus, ScoreRating_b, ScoreRating_bMinus:
-		sd.B++
+		sd.B += count
 	case ScoreRating_cPlus, ScoreRating_c, ScoreRating_cMinus:
-		sd.C++
+		sd.C += count
 	case ScoreRating_dPlus, ScoreRating_d, ScoreRating_dMinus:
-		sd.D++
+		sd.D += count
 	case ScoreRating_failed:
-		sd.F++
+		sd.F += count
 	case ScoreRating_error:
-		sd.Error++
+		sd.Error += count
 	case ScoreRating_unrated:
-		sd.Unrated++
+		sd.Unrated += count
 	}
 }
 
diff --git a/policy/report_test.go b/policy/report_test.go
index 1d97d149..08c7cce7 100644
--- a/policy/report_test.go
+++ b/policy/report_test.go
@@ -40,6 +40,23 @@ func TestScoreDistributionAdd(t *testing.T) {
 	require.Equal(t, uint32(16), scoreDist.GetTotal())
 }
 
+func TestScoreDistributionAddCount(t *testing.T) {
+	scoreDist := &policy.ScoreDistribution{}
+	scoreDist.AddCount(&policy.Score{Value: 100, Type: policy.ScoreType_Result, ScoreCompletion: 100}, 3)
+	scoreDist.AddCount(&policy.Score{Value: 75, Type: policy.ScoreType_Result, ScoreCompletion: 100}, 4)
+	scoreDist.AddCount(&policy.Score{Value: 50, Type: policy.ScoreType_Result, ScoreCompletion: 100}, 5)
+	scoreDist.AddCount(&policy.Score{Type: policy.ScoreType_Error}, 6)
+	scoreDist.AddCount(&policy.Score{Type: policy.ScoreType_Unscored}, 7)
+
+	require.Equal(t, uint32(3), scoreDist.GetA())
+	require.Equal(t, uint32(4), scoreDist.GetB())
+	require.Equal(t, uint32(5), scoreDist.GetC())
+	require.Equal(t, uint32(6), scoreDist.GetError())
+	require.Equal(t, uint32(7), scoreDist.GetUnrated())
+
+	require.Equal(t, uint32(25), scoreDist.GetTotal())
+}
+
 func TestScoreDistributionRemove(t *testing.T) {
 	scoreDist := &policy.ScoreDistribution{
 		Total:   19,