Fix legacy issue with context.done callbacks (#118)

src/utils/handler.spec.ts

@@ -233,4 +233,67 @@ describe("wrap", () => {
     expect(result).toEqual({ statusCode: 204, body: "The callback response" });
     expect(calledOriginalHandler).toBeFalsy();
   });
+
+  it("completes when calling context.done", async () => {
+    const handler: Handler = async (event, context, callback) => {
+      context.done(undefined, { statusCode: 204, body: "The callback response" });
+    };
+
+    let calledOriginalHandler = false;
+
+    const wrappedHandler = wrap(
+      handler,
+      async () => {},
+      async () => {},
+    );
+
+    const result = await wrappedHandler({}, mockContext, () => {
+      calledOriginalHandler = true;
+    });
+
+    expect(result).toEqual({ statusCode: 204, body: "The callback response" });
+    expect(calledOriginalHandler).toBeFalsy();
+  });
+  it("completes when calling context.succeed", async () => {
+    const handler: Handler = async (event, context, callback) => {
+      context.succeed({ statusCode: 204, body: "The callback response" });
+    };
+
+    let calledOriginalHandler = false;
+
+    const wrappedHandler = wrap(
+      handler,
+      async () => {},
+      async () => {},
+    );
+
+    const result = await wrappedHandler({}, mockContext, () => {
+      calledOriginalHandler = true;
+    });
+
+    expect(result).toEqual({ statusCode: 204, body: "The callback response" });
+    expect(calledOriginalHandler).toBeFalsy();
+  });
+
+  it("throws error when calling context.fail", async () => {
+    const handler: Handler = async (event, context, callback) => {
+      context.fail(new Error("Some error"));
+    };
+
+    let calledOriginalHandler = false;
+
+    const wrappedHandler = wrap(
+      handler,
+      async () => {},
+      async () => {},
+    );
+
+    const result = wrappedHandler({}, mockContext, () => {
+      calledOriginalHandler = true;
+    });
+
+    await expect(result).rejects.toEqual(new Error("Some error"));
+
+    expect(calledOriginalHandler).toBeFalsy();
+  });
 });

src/utils/handler.ts

@@ -67,6 +67,9 @@ export function promisifiedHandler<TEvent, TResult>(handler: Handler<TEvent, TRe
     // 2. Returning a value directly from the function using a promise.
 
     let modifiedCallback: Callback<TResult> = () => {};
+    let modifiedLegacyDoneCallback: Callback<TResult> = () => {};
+    let modifiedLegacySucceedCallback: (res: any) => void = () => {};
+    let modifiedLegacyFailCallback: (err: any) => void = () => {};
 
     const callbackProm = new Promise<TResult>((resolve, reject) => {
       modifiedCallback = (err, result) => {
@@ -76,7 +79,28 @@ export function promisifiedHandler<TEvent, TResult>(handler: Handler<TEvent, TRe
           resolve(result);
         }
       };
+      // Legacy done callback finished immediately, and doesn't wait for pending
+      // event loop
+      modifiedLegacyDoneCallback = (err, result) => {
+        context.callbackWaitsForEmptyEventLoop = false;
+        if (err !== undefined && err !== null) {
+          reject(err);
+        } else {
+          resolve(result);
+        }
+      };
+      modifiedLegacySucceedCallback = (result) => {
+        context.callbackWaitsForEmptyEventLoop = false;
+        resolve(result);
+      };
+      modifiedLegacyFailCallback = (err: any) => {
+        context.callbackWaitsForEmptyEventLoop = false;
+        reject(err);
+      };
     });
+    context.done = modifiedLegacyDoneCallback;
+    context.succeed = modifiedLegacySucceedCallback;
+    context.fail = modifiedLegacyFailCallback;
 
     const asyncProm = handler(event, context, modifiedCallback) as Promise<TResult> | undefined;
     let promise: Promise<TResult> = callbackProm;