From 908eca3e87260792aa684a2bec8587993ff32e3f Mon Sep 17 00:00:00 2001
From: Andrea Amorosi <dreamorosi@gmail.com>
Date: Thu, 22 Jun 2023 13:52:05 +0000
Subject: [PATCH] feat(metrics): publish metrics when other middlewares return
 early

---
 packages/metrics/src/middleware/middy.ts      | 15 ++++
 .../tests/unit/middleware/middy.test.ts       | 68 +++++++++++++------
 2 files changed, 62 insertions(+), 21 deletions(-)

diff --git a/packages/metrics/src/middleware/middy.ts b/packages/metrics/src/middleware/middy.ts
index b1a0313214..0da5c04828 100644
--- a/packages/metrics/src/middleware/middy.ts
+++ b/packages/metrics/src/middleware/middy.ts
@@ -1,3 +1,4 @@
+import { METRICS_KEY } from '@aws-lambda-powertools/commons/lib/middleware';
 import type { Metrics } from '../Metrics';
 import type { ExtraOptions } from '../types';
 import type {
@@ -38,6 +39,18 @@ const logMetrics = (
 ): MiddlewareLikeObj => {
   const metricsInstances = target instanceof Array ? target : [target];
 
+  /**
+   * Set the cleanup function to be called in case other middlewares return early.
+   *
+   * @param request - The request object
+   */
+  const setCleanupFunction = (request: MiddyLikeRequest): void => {
+    request.internal = {
+      ...request.internal,
+      [METRICS_KEY]: logMetricsAfterOrError,
+    };
+  };
+
   const logMetricsBefore = async (request: MiddyLikeRequest): Promise<void> => {
     metricsInstances.forEach((metrics: Metrics) => {
       metrics.setFunctionName(request.context.functionName);
@@ -53,6 +66,8 @@ const logMetrics = (
         metrics.captureColdStartMetric();
       }
     });
+
+    setCleanupFunction(request);
   };
 
   const logMetricsAfterOrError = async (): Promise<void> => {
diff --git a/packages/metrics/tests/unit/middleware/middy.test.ts b/packages/metrics/tests/unit/middleware/middy.test.ts
index 1b450bf102..fe2268a78a 100644
--- a/packages/metrics/tests/unit/middleware/middy.test.ts
+++ b/packages/metrics/tests/unit/middleware/middy.test.ts
@@ -11,6 +11,9 @@ import {
 } from '../../../../metrics/src';
 import middy from '@middy/core';
 import { ExtraOptions } from '../../../src/types';
+import { cleanupMiddlewares } from '@aws-lambda-powertools/commons/lib/middleware';
+import { helloworldContext as dummyContext } from '../../../../commons/src/samples/resources/contexts/hello-world';
+import { CustomEvent as dummyEvent } from '../../../../commons/src/samples/resources/events/custom/index';
 
 const consoleSpy = jest.spyOn(console, 'log').mockImplementation();
 const consoleWarnSpy = jest.spyOn(console, 'warn').mockImplementation();
@@ -18,27 +21,6 @@ const mockDate = new Date(1466424490000);
 jest.spyOn(global, 'Date').mockImplementation(() => mockDate);
 
 describe('Middy middleware', () => {
-  const dummyEvent = {
-    key1: 'value1',
-    key2: 'value2',
-    key3: 'value3',
-  };
-  const dummyContext = {
-    callbackWaitsForEmptyEventLoop: true,
-    functionVersion: '$LATEST',
-    functionName: 'foo-bar-function',
-    memoryLimitInMB: '128',
-    logGroupName: '/aws/lambda/foo-bar-function-123456abcdef',
-    logStreamName: '2021/03/09/[$LATEST]abcdef123456abcdef123456abcdef123456',
-    invokedFunctionArn:
-      'arn:aws:lambda:eu-west-1:123456789012:function:Example',
-    awsRequestId: 'c6af9ac6-7b61-11e6-9a41-93e812345678',
-    getRemainingTimeInMillis: () => 1234,
-    done: () => console.log('Done!'),
-    fail: () => console.log('Failed!'),
-    succeed: () => console.log('Succeeded!'),
-  };
-
   beforeEach(() => {
     jest.resetModules();
     jest.clearAllMocks();
@@ -399,6 +381,50 @@ describe('Middy middleware', () => {
         })
       );
     });
+
+    test('when enabled, and another middleware returns early, it still publishes the metrics at the end of the execution', async () => {
+      // Prepare
+      const metrics = new Metrics({
+        namespace: 'serverlessAirline',
+        serviceName: 'orders',
+      });
+      const publishStoredMetricsSpy = jest.spyOn(
+        metrics,
+        'publishStoredMetrics'
+      );
+      const myCustomMiddleware = (): middy.MiddlewareObj => {
+        const before = async (
+          request: middy.Request
+        ): Promise<undefined | string> => {
+          // Return early on the second invocation
+          if (request.event.idx === 1) {
+            // Cleanup Powertools resources
+            await cleanupMiddlewares(request);
+
+            // Then return early
+            return 'foo';
+          }
+        };
+
+        return {
+          before,
+        };
+      };
+      const handler = middy(
+        (_event: typeof dummyEvent & { idx: number }): void => {
+          metrics.addMetric('successfulBooking', MetricUnits.Count, 1);
+        }
+      )
+        .use(logMetrics(metrics))
+        .use(myCustomMiddleware());
+
+      // Act
+      await handler({ ...dummyEvent, idx: 0 }, dummyContext);
+      await handler({ ...dummyEvent, idx: 1 }, dummyContext);
+
+      // Assess
+      expect(publishStoredMetricsSpy).toBeCalledTimes(2);
+    });
   });
   describe('Metrics resolution', () => {
     test('serialized metrics in EMF format should not contain `StorageResolution` as key if `60` is set', async () => {