From fddbfc2a9f1220b8f158d15ae38d3377dcc82d3f Mon Sep 17 00:00:00 2001
From: minnakt <minna.kt@mongodb.com>
Date: Tue, 10 Dec 2024 14:49:36 -0500
Subject: [PATCH] Refactor

---
 model/task/db.go               | 23 +++++------------------
 model/task/task_annotations.go | 30 ++++++++++--------------------
 2 files changed, 15 insertions(+), 38 deletions(-)

diff --git a/model/task/db.go b/model/task/db.go
index 2e4ca0d1c0..96ce334db0 100644
--- a/model/task/db.go
+++ b/model/task/db.go
@@ -2733,18 +2733,6 @@ func enableDisabledTasks(taskIDs []string) error {
 	return err
 }
 
-// SetHasAnnotations sets a task's HasAnnotations flag, indicating
-// that there are annotations with populated IssuesKey for its
-// id / execution pair.
-func SetHasAnnotations(taskId string, execution int) error {
-	err := UpdateOne(
-		ByIdAndExecution(taskId, execution),
-		bson.M{"$set": bson.M{
-			HasAnnotationsKey: true,
-		}})
-	return errors.Wrapf(err, "marking task '%s' as having annotations", taskId)
-}
-
 // IncNumNextTaskDispatches sets the number of times a host has requested this
 // task and execution as its next task.
 func (t *Task) IncNumNextTaskDispatches() error {
@@ -2759,16 +2747,15 @@ func (t *Task) IncNumNextTaskDispatches() error {
 	return nil
 }
 
-// UnsetHasAnnotations unsets a task's HasAnnotations flag, indicating
-// that there are no longer any annotations with populated IssuesKey for its
-// id / execution pair.
-func UnsetHasAnnotations(taskId string, execution int) error {
+// UpdateHasAnnotations updates a task's HasAnnotations flag, indicating if there
+// are any annotations with populated IssuesKey for its id / execution pair.
+func UpdateHasAnnotations(taskId string, execution int, hasAnnotations bool) error {
 	err := UpdateOne(
 		ByIdAndExecution(taskId, execution),
 		bson.M{"$set": bson.M{
-			HasAnnotationsKey: false,
+			HasAnnotationsKey: hasAnnotations,
 		}})
-	return errors.Wrapf(err, "marking task '%s' as having no annotations", taskId)
+	return errors.Wrapf(err, "updating hasAnnotations field for task '%s'", taskId)
 }
 
 type NumExecutionsForIntervalInput struct {
diff --git a/model/task/task_annotations.go b/model/task/task_annotations.go
index e4f60215c3..3b8b54ef5e 100644
--- a/model/task/task_annotations.go
+++ b/model/task/task_annotations.go
@@ -37,7 +37,7 @@ func MoveIssueToSuspectedIssue(taskId string, taskExecution int, issue annotatio
 		return errors.Wrapf(err, "finding and modifying task annotation for execution %d of task '%s'", taskExecution, taskId)
 	}
 	if len(annotation.Issues) == 0 {
-		return UnsetHasAnnotations(taskId, taskExecution)
+		return UpdateHasAnnotations(taskId, taskExecution, false)
 	}
 	return nil
 }
@@ -59,7 +59,7 @@ func MoveSuspectedIssueToIssue(taskId string, taskExecution int, issue annotatio
 	); err != nil {
 		return err
 	}
-	return SetHasAnnotations(taskId, taskExecution)
+	return UpdateHasAnnotations(taskId, taskExecution, true)
 }
 
 // AddIssueToAnnotation adds an issue onto an existing annotation and marks its associated task document
@@ -79,7 +79,7 @@ func AddIssueToAnnotation(taskId string, execution int, issue annotations.IssueL
 	); err != nil {
 		return errors.Wrapf(err, "adding task annotation issue for task '%s'", taskId)
 	}
-	return SetHasAnnotations(taskId, execution)
+	return UpdateHasAnnotations(taskId, execution, true)
 }
 
 // RemoveIssueFromAnnotation removes an issue from an existing annotation, and unsets its
@@ -100,7 +100,7 @@ func RemoveIssueFromAnnotation(taskId string, execution int, issue annotations.I
 		return errors.Wrapf(err, "finding and removing issue for task annotation for execution %d of task '%s'", execution, taskId)
 	}
 	if len(annotation.Issues) == 0 {
-		return UnsetHasAnnotations(taskId, execution)
+		return UpdateHasAnnotations(taskId, execution, false)
 	}
 	return nil
 }
@@ -152,14 +152,9 @@ func UpsertAnnotation(a *annotations.TaskAnnotation, userDisplayName string) err
 	); err != nil {
 		return errors.Wrapf(err, "adding task annotation for task '%s'", a.TaskId)
 	}
-	if a.Issues != nil {
-		if len(a.Issues) == 0 {
-			return UnsetHasAnnotations(a.TaskId, a.TaskExecution)
-		} else {
-			return SetHasAnnotations(a.TaskId, a.TaskExecution)
-		}
-	}
-	return nil
+
+	hasAnnotations := len(a.Issues) > 0
+	return UpdateHasAnnotations(a.TaskId, a.TaskExecution, hasAnnotations)
 }
 
 // PatchAnnotation adds issues onto existing annotations, and marks its associated task document
@@ -199,12 +194,7 @@ func PatchAnnotation(a *annotations.TaskAnnotation, userDisplayName string, upse
 	); err != nil {
 		return errors.Wrapf(err, "updating task annotation for '%s'", a.TaskId)
 	}
-	if a.Issues != nil {
-		if len(a.Issues) == 0 {
-			return UnsetHasAnnotations(a.TaskId, a.TaskExecution)
-		} else {
-			return SetHasAnnotations(a.TaskId, a.TaskExecution)
-		}
-	}
-	return nil
+
+	hasAnnotations := len(a.Issues) > 0
+	return UpdateHasAnnotations(a.TaskId, a.TaskExecution, hasAnnotations)
 }