fix: middleware now stores initial persistent attributes correctly

Author: dreamorosi

packages/logger/src/middleware/middy.ts

@@ -35,9 +35,9 @@ const injectLambdaContext = (target: Logger | Logger[], options?: HandlerOptions
   const persistentAttributes: LogAttributes[] = [];
 
   const injectLambdaContextBefore = async (request: MiddyLikeRequest): Promise<void> => {
-    loggers.forEach((logger: Logger) => {
+    loggers.forEach((logger: Logger, index: number) => {
       if (options && options.clearState === true) {
-        persistentAttributes.push({ ...logger.getPersistentLogAttributes() });
+        persistentAttributes[index] = ({ ...logger.getPersistentLogAttributes() });
       }
       Logger.injectLambdaContextBefore(logger, request.event, request.context, options);
     });

packages/logger/tests/unit/Logger.test.ts

@@ -1227,11 +1227,10 @@ 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 injectLambdaContextAfterOrOnErrorSpy = jest.spyOn(
+        Logger,
+        'injectLambdaContextAfterOrOnError'
+      );
       const logger = new Logger({
         logLevel: 'DEBUG',
       });
@@ -1259,10 +1258,70 @@ describe('Class: Logger', () => {
 
       // 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]);
+      // Here we assert that the logger.info method is called before the cleanup function that should always 
+      // be called ONLY after the handler has returned. If logger.info is called after the cleanup function
+      // it means the decorator is NOT awaiting the handler which would cause the test to fail.
+      expect(consoleSpy.mock.invocationCallOrder[0])
+        .toBeLessThan(injectLambdaContextAfterOrOnErrorSpy.mock.invocationCallOrder[0]);
+
+    });
+
+    test('when logEvent and clearState are both TRUE, and the logger has persistent attributes, any key added in the handler is cleared properly', async () => {
+
+      // Prepare
+      const logger = new Logger({
+        persistentLogAttributes: {
+          version: '1.0.0',
+        }
+      });
+      const consoleSpy = jest.spyOn(logger['console'], 'info').mockImplementation();
+      class LambdaFunction implements LambdaInterface {
+        @logger.injectLambdaContext({ clearState: true, logEvent: true })
+        // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+        // @ts-ignore
+        public async handler(event: { foo: string }, _context: unknown): Promise<unknown> {
+          logger.appendKeys({ foo: event.foo });
+
+          return;
+        }
+      }
+      const lambda = new LambdaFunction();
+      const handler = lambda.handler.bind(lambda);
+
+      // Act
+      await handler({ foo: 'bar' }, {} as Context);
+      await handler({ foo: 'baz' }, {} as Context);
+      await handler({ foo: 'biz' }, {} as Context);
+      await handler({ foo: 'buz' }, {} as Context);
+      await handler({ foo: 'boz' }, {} as Context);
+
+      expect(consoleSpy).toBeCalledTimes(5);
+      for (let i = 1; i === 5; i++) {
+        expect(consoleSpy).toHaveBeenNthCalledWith(
+          i,
+          expect.stringContaining('\"message\":\"Lambda invocation event\"'),
+        );
+        expect(consoleSpy).toHaveBeenNthCalledWith(
+          i,
+          expect.stringContaining('\"version\":\"1.0.0\"'),
+        );
+      }
+      expect(consoleSpy).toHaveBeenNthCalledWith(
+        2,
+        expect.not.stringContaining('\"foo\":\"bar\"')
+      );
+      expect(consoleSpy).toHaveBeenNthCalledWith(
+        3,
+        expect.not.stringContaining('\"foo\":\"baz\"')
+      );
+      expect(consoleSpy).toHaveBeenNthCalledWith(
+        4,
+        expect.not.stringContaining('\"foo\":\"biz\"')
+      );
+      expect(consoleSpy).toHaveBeenNthCalledWith(
+        5,
+        expect.not.stringContaining('\"foo\":\"buz\"')
+      );
 
     });
 

packages/logger/tests/unit/middleware/middy.test.ts

@@ -10,6 +10,7 @@ import { Logger } from './../../../src';
 import middy from '@middy/core';
 import { PowertoolLogFormatter } from '../../../src/formatter';
 import { Console } from 'console';
+import { Context } from 'aws-lambda';
 
 const mockDate = new Date(1466424490000);
 const dateSpy = jest.spyOn(global, 'Date').mockImplementation(() => mockDate);
@@ -491,6 +492,54 @@ describe('Middy middleware', () => {
       }));
     });
 
+    test('when logEvent and clearState are both TRUE, and the logger has persistent attributes, any key added in the handler is cleared properly', async () => {
+
+      const logger = new Logger({
+        persistentLogAttributes: {
+          version: '1.0.0',
+        }
+      });
+      const consoleSpy = jest.spyOn(logger['console'], 'info').mockImplementation();
+      const handler = middy(async (event: { foo: string }, _context: Context) => {
+        logger.appendKeys({ foo: event.foo });
+      }).use(injectLambdaContext(logger, { logEvent: true, clearState: true }));
+
+      await handler({ foo: 'bar' }, {} as Context);
+      await handler({ foo: 'baz' }, {} as Context);
+      await handler({ foo: 'biz' }, {} as Context);
+      await handler({ foo: 'buz' }, {} as Context);
+      await handler({ foo: 'boz' }, {} as Context);
+
+      expect(consoleSpy).toBeCalledTimes(5);
+      for (let i = 1; i === 5; i++) {
+        expect(consoleSpy).toHaveBeenNthCalledWith(
+          i,
+          expect.stringContaining('\"message\":\"Lambda invocation event\"'),
+        );
+        expect(consoleSpy).toHaveBeenNthCalledWith(
+          i,
+          expect.stringContaining('\"version\":\"1.0.0\"'),
+        );
+      }
+      expect(consoleSpy).toHaveBeenNthCalledWith(
+        2,
+        expect.not.stringContaining('\"foo\":\"bar\"')
+      );
+      expect(consoleSpy).toHaveBeenNthCalledWith(
+        3,
+        expect.not.stringContaining('\"foo\":\"baz\"')
+      );
+      expect(consoleSpy).toHaveBeenNthCalledWith(
+        4,
+        expect.not.stringContaining('\"foo\":\"biz\"')
+      );
+      expect(consoleSpy).toHaveBeenNthCalledWith(
+        5,
+        expect.not.stringContaining('\"foo\":\"buz\"')
+      );
+
+    });
+
   });
 
 });