Incorporate feedback

Author: ConnorKirk

Diff for: packages/logger/src/middleware/middy.ts

@@ -77,11 +77,11 @@ const injectLambdaContext = (
   const setCleanupFunction = (request: MiddyLikeRequest): void => {
     request.internal = {
       ...request.internal,
-      [LOGGER_KEY]: onAfter,
+      [LOGGER_KEY]: after,
     };
   };
 
-  const onBefore = async (request: MiddyLikeRequest): Promise<void> => {
+  const before = async (request: MiddyLikeRequest): Promise<void> => {
     for (const logger of loggers) {
       if (isResetStateEnabled) {
         setCleanupFunction(request);
@@ -98,18 +98,17 @@ const injectLambdaContext = (
     }
   };
 
-  const onAfter = async (): Promise<void> => {
-    if (isResetStateEnabled) {
-      for (const logger of loggers) {
-        logger.resetKeys();
-      }
-    }
-
+  const after = async (): Promise<void> => {
     for (const logger of loggers) {
       logger.clearBuffer();
+
+      if (isResetStateEnabled) {
+        logger.resetKeys();
+      }
     }
   };
-  const onError = async ({ error }): Promise<void> => {
+
+  const onError = async ({ error }: { error: unknown }): Promise<void> => {
     if (options?.flushBufferOnUncaughtError) {
       for (const logger of loggers) {
         logger.flushBuffer();
@@ -122,8 +121,8 @@ const injectLambdaContext = (
   };
 
   return {
-    before: onBefore,
-    after: onAfter,
+    before,
+    after,
     onError,
   };
 };

Diff for: packages/logger/tests/unit/logBuffer.test.ts

@@ -243,74 +243,62 @@ describe('Buffer logs', () => {
     expect(console.error).toBeCalledTimes(1);
   });
 
-  it('flushes the buffer when an uncaught error is thrown', async () => {
-    // Prepare
-    const logger = new Logger({ logBufferOptions: { enabled: true } });
-    class TestClass {
-      @logger.injectLambdaContext({ flushBufferOnUncaughtError: true })
-      async handler(_event: unknown, _context: Context) {
-        logger.debug('This is a log message');
-        logger.info('This is an info message');
-        throw new Error('This is an error');
-      }
+  it.each([
+    {
+      handlerFactory: (logger: Logger) =>
+        middy()
+          .use(
+            injectLambdaContext(logger, { flushBufferOnUncaughtError: true })
+          )
+          .handler(async () => {
+            logger.debug('This is a log message');
+            logger.info('This is an info message');
+            throw new Error('This is an error');
+          }),
+      case: 'middleware',
+    },
+    {
+      handlerFactory: (logger: Logger) => {
+        class TestClass {
+          @logger.injectLambdaContext({ flushBufferOnUncaughtError: true })
+          async handler(_event: unknown, _context: Context) {
+            logger.debug('This is a log message');
+            logger.info('This is an info message');
+            throw new Error('This is an error');
+          }
+        }
+        const lambda = new TestClass();
+        return lambda.handler.bind(lambda);
+      },
+      case: 'decorator',
+    },
+  ])(
+    'flushes the buffer when an uncaught error is thrown ($case)',
+    async ({ handlerFactory }) => {
+      // Prepare
+      const logger = new Logger({ logBufferOptions: { enabled: true } });
+      const handler = handlerFactory(logger);
+
+      // Act & Assess
+      await expect(() =>
+        handler(
+          {
+            foo: 'bar',
+          },
+          context
+        )
+      ).rejects.toThrow(new Error('This is an error'));
+      expect(console.debug).toBeCalledTimes(1);
+      expect(console.info).toBeCalledTimes(1);
+      expect(console.error).toHaveLogged(
+        expect.objectContaining({
+          message: UncaughtErrorLogMessage,
+        })
+      );
+      // If debug is called after info, it means it was buffered and then flushed
+      expect(console.debug).toHaveBeenCalledAfter(console.info as Mock);
+      // If error is called after debug, it means the buffer was flushed before the error log
+      expect(console.debug).toHaveBeenCalledBefore(console.error as Mock);
     }
-    const lambda = new TestClass();
-    const handler = lambda.handler.bind(lambda);
-
-    // Act & Assess
-    await expect(() =>
-      handler(
-        {
-          foo: 'bar',
-        },
-        context
-      )
-    ).rejects.toThrow(new Error('This is an error'));
-    expect(console.debug).toBeCalledTimes(1);
-    expect(console.info).toBeCalledTimes(1);
-    expect(console.error).toHaveLogged(
-      expect.objectContaining({
-        message: UncaughtErrorLogMessage,
-      })
-    );
-    // If debug is called after info, it means it was buffered and then flushed
-    expect(console.debug).toHaveBeenCalledAfter(console.info as Mock);
-    // If error is called after debug, it means the buffer was flushed before the error log
-    expect(console.debug).toHaveBeenCalledBefore(console.error as Mock);
-  });
-  it('flushes the buffer when an uncaught error is thrown in a middy', async () => {
-    // Prepare
-    const logger = new Logger({ logBufferOptions: { enabled: true } });
-
-    const handlerFn = async (_event: unknown, _context: Context) => {
-      logger.debug('This is a log message');
-      logger.info('This is an info message');
-      throw new Error('This is an error');
-    };
-
-    const handler = middy()
-      .use(injectLambdaContext(logger, { flushBufferOnUncaughtError: true }))
-      .handler(handlerFn);
-
-    // Act & Assess
-    await expect(() =>
-      handler(
-        {
-          foo: 'bar',
-        },
-        context
-      )
-    ).rejects.toThrow(new Error('This is an error'));
-    expect(console.debug).toBeCalledTimes(1);
-    expect(console.info).toBeCalledTimes(1);
-    expect(console.error).toHaveLogged(
-      expect.objectContaining({
-        message: UncaughtErrorLogMessage,
-      })
-    );
-    // If debug is called after info, it means it was buffered and then flushed
-    expect(console.debug).toHaveBeenCalledAfter(console.info as Mock);
-    // If error is called after debug, it means the buffer was flushed before the error log
-    expect(console.debug).toHaveBeenCalledBefore(console.error as Mock);
-  });
+  );
 });