diff --git a/packages/logger/src/Logger.ts b/packages/logger/src/Logger.ts index 2bbdb3fcc1..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. */ @@ -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 --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(_event: unknown, _context: Context, _callback: Callback): void | Promise { + await this.dummyMethod(); + logger.info('This is a DEBUG log'); + + return; + } + + private async dummyMethod(): Promise { + 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', () => {