support to update task status via Builder

Signed-off-by: Varun Bansal <bansvaru@amazon.com>
opensearch-project · Aug 5, 2024 · 3a84e95 · 3a84e95
1 parent 6470881
commit 3a84e95
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/libs/task-commons/src/main/java/org/opensearch/task/commons/task/Task.java b/libs/task-commons/src/main/java/org/opensearch/task/commons/task/Task.java
@@ -198,7 +198,7 @@ public static class Builder {
         /**
          * Depicts latest state of the Task
          */
-        private final TaskStatus taskStatus;
+        private TaskStatus taskStatus;
 
         /**
          * Various params to used for Task execution
@@ -323,6 +323,14 @@ public void lastHeartbeatAt(long lastHeartbeatAt) {
             this.lastHeartbeatAt = lastHeartbeatAt;
         }
 
+        /**
+         * Update the Task Status
+         * @param taskStatus {@link TaskStatus} - current status of the Task
+         */
+        public void taskStatus(TaskStatus taskStatus) {
+            this.taskStatus = taskStatus;
+        }
+
         /**
          * Set Task Assigned Node
          * @param node Worker Node on which the Task is to be executed

diff --git a/libs/task-commons/src/test/java/org/opensearch/task/commons/task/TaskTests.java b/libs/task-commons/src/test/java/org/opensearch/task/commons/task/TaskTests.java
@@ -107,11 +107,12 @@ public void testBuilderFromAttributes() {
         builder.completedAt(createdAt + 3000);
         builder.lastHeartbeatAt(createdAt + 2500);
         builder.assignedNode(WorkerNode.createWorkerNode("node1", "nodeip", "nodename"));
+        builder.taskStatus(TaskStatus.ACTIVE);
 
         Task task = builder.build();
 
         assertEquals(taskId, task.getTaskId());
-        assertEquals(taskStatus, task.getTaskStatus());
+        assertEquals(TaskStatus.ACTIVE, task.getTaskStatus());
         assertEquals(params, task.getParams());
         assertEquals(taskType, task.getTaskType());
         assertEquals(createdAt, task.getCreatedAt());