waitForWorflowCompletion

Author: Shailesh Jagannath Padave

src/core/__test__/utils/test-util.ts

@@ -2,6 +2,7 @@ import * as fs from 'fs';
 import * as path from 'path';
 import {WorkflowDef} from "../../../common";
 import {MetadataClient} from "../../metadataClient";
+import {WorkflowExecutor} from "../../executor";
 
 export class TestUtil {
     private static metadataClient: MetadataClient;
@@ -35,6 +36,39 @@ export class TestUtil {
         }
     }
 
+    // Helper function to wait for workflow completion
+    public static async waitForWorkflowCompletion(
+        executor: WorkflowExecutor,
+        workflowId: string,
+        maxWaitMs: number = 300000, // 5 minutes default
+        pollIntervalMs: number = 100 // 100ms default
+    ) {
+        const startTime = Date.now();
+
+        while (Date.now() - startTime < maxWaitMs) {
+            try {
+                const workflowStatus = await executor.getWorkflow(workflowId, true);
+
+                // Check if workflow is in a terminal state
+                if (['COMPLETED', 'FAILED', 'TERMINATED', 'TIMED_OUT'].includes(workflowStatus.status!)) {
+                    console.debug(`Workflow ${workflowId} reached terminal state: ${workflowStatus.status}`);
+                    return workflowStatus;
+                }
+
+                console.debug(`Workflow ${workflowId} status: ${workflowStatus.status}, waiting...`);
+
+                // Wait before next poll
+                await new Promise(resolve => setTimeout(resolve, pollIntervalMs));
+
+            } catch (error) {
+                console.warn(`Error checking workflow status for ${workflowId}:`, error);
+                await new Promise(resolve => setTimeout(resolve, pollIntervalMs));
+            }
+        }
+
+        throw new Error(`Workflow ${workflowId} did not complete within ${maxWaitMs}ms`);
+    }
+
     /**
      * Unregister a workflow
      */

src/task/__tests__/TaskManager.test.ts

@@ -3,6 +3,7 @@ import { simpleTask, WorkflowExecutor } from "../../core";
 import { orkesConductorClient } from "../../orkes";
 import { TaskManager, ConductorWorker } from "../index";
 import { mockLogger } from "./mockLogger";
+import {TestUtil} from "../../core/__test__/utils/test-util";
 
 
 const BASE_TIME = 500;
@@ -169,9 +170,9 @@ describe("TaskManager", () => {
     // decrease speed again
     manager.updatePollingOptions({ pollInterval: BASE_TIME, concurrency: 1 });
 
-    const workflowStatus = await executor.getWorkflow(executionId!, true);
+    const workflowStatus = TestUtil.waitForWorkflowCompletion(executor, executionId!, 30000);
 
-    expect(workflowStatus.status).toEqual("COMPLETED");
+    expect(workflowStatus).toEqual("COMPLETED");
     await manager.stopPolling();
 
     expect(manager.isPolling).toBeFalsy();