diff --git a/packages/logger/src/middleware/middy.ts b/packages/logger/src/middleware/middy.ts index eeaef4fd49..fe36bfaf18 100644 --- a/packages/logger/src/middleware/middy.ts +++ b/packages/logger/src/middleware/middy.ts @@ -123,6 +123,10 @@ const injectLambdaContext = ( } else { logger.clearBuffer(); } + + if (isResetStateEnabled) { + logger.resetKeys(); + } } }; diff --git a/packages/logger/tests/unit/injectLambdaContext.test.ts b/packages/logger/tests/unit/injectLambdaContext.test.ts index b3861b74ad..03dec2cb66 100644 --- a/packages/logger/tests/unit/injectLambdaContext.test.ts +++ b/packages/logger/tests/unit/injectLambdaContext.test.ts @@ -415,4 +415,121 @@ describe('Inject Lambda Context', () => { }) ); }); + + it.each([ + { + case: 'middleware', + getHandler: (logger: Logger) => + middy(async (event: { id: number }) => { + logger.info('Processing event'); + logger.appendKeys({ id: event.id }); + throw new Error('Test error'); + }).use( + injectLambdaContext(logger, { + resetKeys: true, + }) + ), + }, + { + case: 'decorator', + getHandler: (logger: Logger) => { + class Lambda { + @logger.injectLambdaContext({ + resetKeys: true, + }) + public async handler( + event: { id: number }, + _context: Context + ): Promise { + logger.info('Processing event'); + logger.appendKeys({ id: event.id }); + throw new Error('Test error'); + } + } + const lambda = new Lambda(); + return lambda.handler.bind(lambda); + }, + }, + ])( + 'resets keys when the handler throws an error $case', + async ({ getHandler }) => { + // Prepare + const logger = new Logger(); + const handler = getHandler(logger); + + // Act & Assess + await expect(handler({ id: 1 }, context)).rejects.toThrow('Test error'); + await expect(handler({ id: 2 }, context)).rejects.toThrow('Test error'); + expect(console.info).toHaveBeenCalledTimes(2); + expect(console.info).toHaveLoggedNth( + 1, + expect.objectContaining({ + message: 'Processing event', + }) + ); + expect(console.info).toHaveLoggedNth( + 2, + expect.not.objectContaining({ + message: 'Processing event', + id: 1, + }) + ); + } + ); + + it.each([ + { + case: 'middleware', + getHandler: (logger: Logger) => + middy(async (event: { id: number }) => { + logger.info('Processing event'); + logger.appendKeys({ id: event.id }); + return true; + }).use( + injectLambdaContext(logger, { + resetKeys: true, + }) + ), + }, + { + case: 'decorator', + getHandler: (logger: Logger) => { + class Lambda { + @logger.injectLambdaContext({ + resetKeys: true, + }) + public async handler(event: { id: number }, _context: Context) { + logger.info('Processing event'); + logger.appendKeys({ id: event.id }); + return true; + } + } + const lambda = new Lambda(); + return lambda.handler.bind(lambda); + }, + }, + ])('resets keys when the handler returns $case', async ({ getHandler }) => { + // Prepare + const logger = new Logger(); + const handler = getHandler(logger); + // Act + await handler({ id: 1 }, context); + await handler({ id: 2 }, context); + + // Assess + expect(console.info).toHaveBeenCalledTimes(2); + expect(console.info).toHaveLoggedNth( + 1, + expect.objectContaining({ + message: 'Processing event', + }) + ); + expect(console.info).toHaveLoggedNth( + 2, + expect.not.objectContaining({ + message: 'Processing event', + id: 1, + }) + ); + }); });