Merge pull request #21 from nibble-4bits/feature/wait-state-pause-ove…

…rride Feature/wait state pause override
nibble-4bits · Dec 14, 2022 · 8d4d40b · 8d4d40b
2 parents 3826462 + 012db3f
commit 8d4d40b
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 6 deletions.
diff --git a/__tests__/StateMachine.test.ts b/__tests__/StateMachine.test.ts
@@ -852,7 +852,7 @@ describe('State Machine', () => {
       const stateMachine = new StateMachine(definition, { checkArn: false });
       const result = await stateMachine.run(input, {
         overrides: {
-          taskResourceLocalHandler: {
+          taskResourceLocalHandlers: {
             TaskState: localHandlerFn,
           },
         },
@@ -1215,6 +1215,37 @@ describe('State Machine', () => {
 
       expect(mockSleepFunction).toHaveBeenCalledWith(20700000);
     });
+
+    test('should pause execution for the specified amount of milliseconds if wait time override option is set', async () => {
+      const definition: StateMachineDefinition = {
+        StartAt: 'PassState',
+        States: {
+          PassState: {
+            Type: 'Pass',
+            Parameters: { waitUntil: '2022-12-05T05:45:00Z' },
+            Next: 'WaitState',
+          },
+          WaitState: {
+            Type: 'Wait',
+            TimestampPath: '$.waitUntil',
+            End: true,
+          },
+        },
+      };
+      const input = {};
+
+      const stateMachine = new StateMachine(definition);
+      await stateMachine.run(input, {
+        overrides: {
+          waitTimeOverrides: {
+            WaitState: 1500,
+          },
+        },
+      });
+
+      expect(mockSleepFunction).toHaveBeenCalledTimes(1);
+      expect(mockSleepFunction).toHaveBeenCalledWith(1500);
+    });
   });
 
   describe('Choice State', () => {

diff --git a/src/StateMachine.ts b/src/StateMachine.ts
@@ -262,8 +262,8 @@ export class StateMachine {
 
     try {
       // If local override for task resource is defined, use that
-      if (options?.overrides?.taskResourceLocalHandler?.[this.currStateName]) {
-        const overrideFn = options?.overrides?.taskResourceLocalHandler?.[this.currStateName];
+      if (options?.overrides?.taskResourceLocalHandlers?.[this.currStateName]) {
+        const overrideFn = options?.overrides?.taskResourceLocalHandlers?.[this.currStateName];
         const result = await overrideFn(this.currInput);
         this.currResult = result;
         return;
@@ -351,9 +351,17 @@ export class StateMachine {
    * Pauses the state machine execution for a certain amount of time
    * based on one of the `Seconds`, `Timestamp`, `SecondsPath` or `TimestampPath` fields.
    */
-  // eslint-disable-next-line @typescript-eslint/no-unused-vars
-  private async handleWaitState(_options?: RunOptions): Promise<void> {
+  private async handleWaitState(options?: RunOptions): Promise<void> {
     const state = this.currState as WaitState;
+    const waitTimeOverrideOption = options?.overrides?.waitTimeOverrides?.[this.currStateName];
+
+    if (typeof waitTimeOverrideOption === 'number') {
+      // If the wait time override is set and is a number, sleep for the specified number of milliseconds
+      // Else, if set to `true`, simply finish the state without sleeping at all.
+      await sleep(waitTimeOverrideOption);
+      this.currResult = this.currInput;
+      return;
+    }
 
     if (state.Seconds) {
       await sleep(state.Seconds * 1000);

diff --git a/src/typings/StateMachineImplementation.ts b/src/typings/StateMachineImplementation.ts
@@ -4,8 +4,13 @@ type TaskStateResourceLocalHandler = {
   [taskStateName: string]: (...args: any) => any;
 };
 
+type WaitStateTimeOverride = {
+  [waitStateName: string]: number;
+};
+
 interface Overrides {
-  taskResourceLocalHandler?: TaskStateResourceLocalHandler;
+  taskResourceLocalHandlers?: TaskStateResourceLocalHandler;
+  waitTimeOverrides?: WaitStateTimeOverride;
 }
 
 export interface RunOptions {