From 65c5509d40f062c78b66c6570c564b686b2f81b8 Mon Sep 17 00:00:00 2001
From: Viren Baraiya <virenx@gmail.com>
Date: Fri, 10 Feb 2023 00:45:53 -0800
Subject: [PATCH] Update WorkflowTestRequest.java

---
 .../client/model/WorkflowTestRequest.java     | 25 ++++++++++++++++---
 1 file changed, 22 insertions(+), 3 deletions(-)

diff --git a/src/main/java/io/orkes/conductor/client/model/WorkflowTestRequest.java b/src/main/java/io/orkes/conductor/client/model/WorkflowTestRequest.java
index 24a8e0e4..cde718c8 100644
--- a/src/main/java/io/orkes/conductor/client/model/WorkflowTestRequest.java
+++ b/src/main/java/io/orkes/conductor/client/model/WorkflowTestRequest.java
@@ -1,5 +1,6 @@
 package io.orkes.conductor.client.model;
 
+import com.netflix.conductor.common.metadata.tasks.Task;
 import com.netflix.conductor.common.metadata.tasks.TaskResult;
 import com.netflix.conductor.common.metadata.workflow.StartWorkflowRequest;
 import lombok.AllArgsConstructor;
@@ -8,6 +9,7 @@
 import lombok.NoArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Queue;
@@ -17,22 +19,39 @@
 public class WorkflowTestRequest extends StartWorkflowRequest {
 
     //Map of task reference name to mock output for the task
-    private Map<String, List<TaskMock>> taskRefToMockOutput;
+    private Map<String, List<TaskMock>> taskRefToMockOutput = new HashMap<>();
+
+    //If there are sub-workflows inside the workflow
+    //The map of task reference name to the mock for the sub-workflow
+    private Map<String, WorkflowTestRequest> subWorkflowTestRequest = new HashMap<>();
 
     @Data
     @NoArgsConstructor
     @AllArgsConstructor
     public static class TaskMock {
         private TaskResult.Status status = TaskResult.Status.COMPLETED;
-        Map<String, Object> output;
+        private Map<String, Object> output = new HashMap<>();
         //Time in millis for the execution of the task.
         //Set this value to view the execution timeline view and also trigger any timeouts associated with the tasks
-        long executionTime;
+        private long executionTime;
+
+        private long queueWaitTime;
 
         public TaskMock(TaskResult.Status status, Map<String, Object> output) {
             this.status = status;
             this.output = output;
         }
+
+        public TaskMock(Task.Status status, Map<String, Object> output, long executionTime, long queueWaitTime) {
+            try {
+                this.status = TaskResult.Status.valueOf(status.name());
+            } catch (IllegalArgumentException invalidValue) {
+                this.status = TaskResult.Status.IN_PROGRESS;
+            }
+            this.output = output;
+            this.executionTime = executionTime;
+            this.queueWaitTime = queueWaitTime;
+        }
     }
 
 }