fix(logger): wait for decorated method return before clearing out state (#1087)

* fix: await decorated method

* chore: housekeeping

* chore: add unit test case

Author: dreamorosi

Diff for: 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.
        */
@@ -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 {

Diff for: 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', () => {