feat(testing-sdk): add ability to wait for submitted status for waitForCallback (#299)

*Issue #, if available:*

#187

*Description of changes:*

Currently the statuses for operation.waitForData only has STARTED and
COMPLETED.

Submitted is useful for waiting for after a waitForCallback has been
submitted instead of just created. This allows the submitter to run to
completion instead of having the callback be instantly completed while
the submitter has not yet completed.

By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice.

Author: anthonyting

packages/aws-durable-execution-sdk-js-examples/src/examples/wait-for-callback/multiple-invocations/wait-for-callback-multiple-invocations.test.ts

@@ -20,20 +20,14 @@ createTests({
         payload: { test: "multiple-invocations" },
       });
 
-      // Wait for first callback and complete it
-      await firstCallbackOp.waitForData(WaitingOperationStatus.STARTED);
-
-      // Give time for callback to be submitted (TODO: https://github.com/aws/aws-durable-execution-sdk-js/issues/187)
-      await new Promise((resolve) => setTimeout(resolve, 1000));
+      // Wait for first callback to be submitted, then complete it
+      await firstCallbackOp.waitForData(WaitingOperationStatus.SUBMITTED);
 
       const firstCallbackResult = JSON.stringify({ step: 1 });
       await firstCallbackOp.sendCallbackSuccess(firstCallbackResult);
 
-      // Wait for second callback and complete it
-      await secondCallbackOp.waitForData(WaitingOperationStatus.STARTED);
-
-      // Give time for callback to be submitted (TODO: https://github.com/aws/aws-durable-execution-sdk-js/issues/187)
-      await new Promise((resolve) => setTimeout(resolve, 1000));
+      // Wait for second callback to be submitted, then complete it
+      await secondCallbackOp.waitForData(WaitingOperationStatus.SUBMITTED);
 
       const secondCallbackResult = JSON.stringify({ step: 2 });
       await secondCallbackOp.sendCallbackSuccess(secondCallbackResult);

packages/aws-durable-execution-sdk-js-testing/src/test-runner/cloud/cloud-durable-test-runner.ts

@@ -49,8 +49,8 @@ export class CloudDurableTestRunner<ResultType>
   private readonly functionArn: string;
   private readonly client: LambdaClient;
   private readonly formatter = new ResultFormatter<ResultType>();
-  private readonly waitManager = new OperationWaitManager();
   private indexedOperations = new IndexedOperations([]);
+  private waitManager = new OperationWaitManager(this.indexedOperations);
   private operationStorage: OperationStorage;
   private readonly apiClient: DurableApiClient;
   private readonly config: CloudDurableTestRunnerConfigInternal;
@@ -190,6 +190,7 @@ export class CloudDurableTestRunner<ResultType>
 
   reset() {
     this.indexedOperations = new IndexedOperations([]);
+    this.waitManager = new OperationWaitManager(this.indexedOperations);
     this.operationStorage = new OperationStorage(
       this.waitManager,
       this.indexedOperations,

packages/aws-durable-execution-sdk-js-testing/src/test-runner/durable-test-runner.ts

@@ -13,7 +13,19 @@ import {
 import { OperationWithData } from "./common/operations/operation-with-data";
 
 export enum WaitingOperationStatus {
+  /**
+   * Fires when the operation starts.
+   */
   STARTED = "STARTED",
+  /**
+   * Submitted is the same as COMPLETED, except for the case where the operation is waitForCallback.
+   * In that case, SUBMITTED will fire when the waitForCallback submitter is completed.
+   */
+  SUBMITTED = "SUBMITTED",
+  /**
+   * Fires when the operation is completed. This includes a status of CANCELLED, FAILED, STOPPED,
+   * SUCCEEDED, or TIMED_OUT.
+   */
   COMPLETED = "COMPLETED",
 }
 

packages/aws-durable-execution-sdk-js-testing/src/test-runner/local/__tests__/local-durable-test-runner.test.ts

@@ -8,6 +8,7 @@ import { InvocationStatus } from "@aws/durable-execution-sdk-js";
 import { CheckpointServerWorkerManager } from "../checkpoint-server-worker-manager";
 import { CheckpointApiClient } from "../api-client/checkpoint-api-client";
 import { FunctionStorage } from "../operations/function-storage";
+import { IndexedOperations } from "../../common/indexed-operations";
 import { install } from "@sinonjs/fake-timers";
 
 jest.mock("../test-execution-orchestrator");
@@ -122,7 +123,9 @@ describe("LocalDurableTestRunner", () => {
       });
 
       expect(runner).toBeDefined();
-      expect(OperationWaitManager).toHaveBeenCalledWith();
+      expect(OperationWaitManager).toHaveBeenCalledWith(
+        expect.any(IndexedOperations),
+      );
       expect(LocalOperationStorage).toHaveBeenCalledWith(
         mockWaitManager,
         expect.any(Object),

packages/aws-durable-execution-sdk-js-testing/src/test-runner/local/local-durable-test-runner.ts

@@ -95,7 +95,7 @@ export class LocalDurableTestRunner<ResultType>
   implements DurableTestRunner<OperationWithData, ResultType>
 {
   private operationStorage: LocalOperationStorage;
-  private readonly waitManager: OperationWaitManager;
+  private waitManager: OperationWaitManager;
   private readonly resultFormatter: ResultFormatter<ResultType>;
   private operationIndex: IndexedOperations;
   static skipTime = false;
@@ -111,8 +111,8 @@ export class LocalDurableTestRunner<ResultType>
    * @param params.handlerFunction The durable function handler to execute
    */
   constructor({ handlerFunction }: LocalDurableTestRunnerParameters) {
-    this.waitManager = new OperationWaitManager();
     this.operationIndex = new IndexedOperations([]);
+    this.waitManager = new OperationWaitManager(this.operationIndex);
     this.resultFormatter = new ResultFormatter<ResultType>();
 
     this.handlerFunction = handlerFunction;
@@ -307,6 +307,7 @@ export class LocalDurableTestRunner<ResultType>
 
   reset() {
     this.operationIndex = new IndexedOperations([]);
+    this.waitManager = new OperationWaitManager(this.operationIndex);
     this.operationStorage = new LocalOperationStorage(
       this.waitManager,
       this.operationIndex,