From f6deb001d213e7f192e9766e808b52ffa164f4ea Mon Sep 17 00:00:00 2001
From: Andrea Amorosi <dreamorosi@gmail.com>
Date: Sun, 11 Sep 2022 16:29:40 +0200
Subject: [PATCH 1/3] fix: await decorated method

---
 packages/logger/src/Logger.ts | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/packages/logger/src/Logger.ts b/packages/logger/src/Logger.ts
index 2bbdb3fcc1..33d7dc0eb4 100644
--- a/packages/logger/src/Logger.ts
+++ b/packages/logger/src/Logger.ts
@@ -285,7 +285,7 @@ class Logger extends Utility implements ClassThatLogs {
       const loggerRef = this;
       // Use a function() {} instead of an () => {} arrow function so that we can
       // access `myClass` as `this` in a decorated `myClass.myMethod()`.
-      descriptor.value = (function (this: Handler, event, context, callback) {
+      descriptor.value = (async function (this: Handler, event, context, callback) {
 
         let initialPersistentAttributes = {};
         if (options && options.clearState === true) {
@@ -297,7 +297,7 @@ class Logger extends Utility implements ClassThatLogs {
         /* eslint-disable  @typescript-eslint/no-non-null-assertion */
         let result: unknown;
         try {
-          result = originalMethod!.apply(this, [ event, context, callback ]);
+          result = await originalMethod!.apply(this, [ event, context, callback ]);
         } catch (error) {
           throw error;
         } finally {

From e1f8790742d7758de56d336a69075e123653cee7 Mon Sep 17 00:00:00 2001
From: Andrea Amorosi <dreamorosi@gmail.com>
Date: Mon, 19 Sep 2022 11:55:57 +0300
Subject: [PATCH 2/3] chore: housekeeping

---
 packages/logger/src/Logger.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/packages/logger/src/Logger.ts b/packages/logger/src/Logger.ts
index 33d7dc0eb4..805e7a5d1a 100644
--- a/packages/logger/src/Logger.ts
+++ b/packages/logger/src/Logger.ts
@@ -275,7 +275,7 @@ class Logger extends Utility implements ClassThatLogs {
    * @returns {HandlerMethodDecorator}
    */
   public injectLambdaContext(options?: HandlerOptions): HandlerMethodDecorator {
-    return (target, _propertyKey, descriptor) => {
+    return (_target, _propertyKey, descriptor) => {
       /**
        * The descriptor.value is the method this decorator decorates, it cannot be undefined.
        */

From a0e07fb340f4f7fa3aa9d1aa9290bbe653248404 Mon Sep 17 00:00:00 2001
From: Andrea Amorosi <dreamorosi@gmail.com>
Date: Mon, 19 Sep 2022 12:34:00 +0300
Subject: [PATCH 3/3] chore: add unit test case

---
 packages/logger/tests/unit/Logger.test.ts | 42 +++++++++++++++++++++++
 1 file changed, 42 insertions(+)

diff --git a/packages/logger/tests/unit/Logger.test.ts b/packages/logger/tests/unit/Logger.test.ts
index 14efa09204..17b7969230 100644
--- a/packages/logger/tests/unit/Logger.test.ts
+++ b/packages/logger/tests/unit/Logger.test.ts
@@ -1166,6 +1166,48 @@ describe('Class: Logger', () => {
 
     });
 
+    test('when used as decorator on an async method, the method is awaited correctly', async () => {
+
+      // Prepare
+      const injectLambdaContextAfterOrOnErrorMock = jest.fn().mockReturnValue('worked');
+      // Temporarily override the cleanup static method so that we can "spy" on it.
+      // This method is always called after the handler has returned in the decorator
+      // implementation.
+      Logger.injectLambdaContextAfterOrOnError = injectLambdaContextAfterOrOnErrorMock;
+      const logger = new Logger({
+        logLevel: 'DEBUG',
+      });
+      const consoleSpy = jest.spyOn(logger['console'], 'info').mockImplementation();
+      class LambdaFunction implements LambdaInterface {
+        @logger.injectLambdaContext()
+        // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+        // @ts-ignore
+        public async handler<TResult>(_event: unknown, _context: Context, _callback: Callback<TResult>): void | Promise<TResult> {
+          await this.dummyMethod();
+          logger.info('This is a DEBUG log');
+
+          return;
+        }
+
+        private async dummyMethod(): Promise<void> {
+          return;
+        }
+      }
+
+      // Act
+      const lambda = new LambdaFunction();
+      const handler = lambda.handler.bind(lambda);
+      await handler({}, dummyContext, () => console.log('Lambda invoked!'));
+
+      // Assess
+      expect(consoleSpy).toBeCalledTimes(1);
+      // Here we assert that the logger.info method is called before the cleanup function that should awlays 
+      // be called after the handler has returned. If logger.info is called after it means the decorator is
+      // NOT awaiting the handler which would cause the test to fail.
+      expect(consoleSpy.mock.invocationCallOrder[0]).toBeLessThan(injectLambdaContextAfterOrOnErrorMock.mock.invocationCallOrder[0]);
+
+    });
+
   });
 
   describe('Method: refreshSampleRateCalculation', () => {