Update how a request handler method is identified (#1058)

humanzz · web-flow · commit e72e82c53c4c · 2023-02-21T10:57:45.000+01:00
- Stop relying on handler method name
- Leverage placedOnRequestHandler and placedOnStreamHandler instead
diff --git a/powertools-core/src/main/java/software/amazon/lambda/powertools/core/internal/LambdaHandlerProcessor.java b/powertools-core/src/main/java/software/amazon/lambda/powertools/core/internal/LambdaHandlerProcessor.java
@@ -37,9 +37,7 @@ private LambdaHandlerProcessor() {
     }
 
     public static boolean isHandlerMethod(final ProceedingJoinPoint pjp) {
-        return "handleRequest".equals(pjp.getSignature().getName()) ||
-                // https://docs.aws.amazon.com/codeguru/latest/profiler-ug/lambda-custom.html
-                "requestHandler".equals(pjp.getSignature().getName());
+        return placedOnRequestHandler(pjp) || placedOnStreamHandler(pjp);
     }
 
     public static boolean placedOnRequestHandler(final ProceedingJoinPoint pjp) {
diff --git a/powertools-core/src/test/java/software/amazon/lambda/powertools/core/internal/LambdaHandlerProcessorTest.java b/powertools-core/src/test/java/software/amazon/lambda/powertools/core/internal/LambdaHandlerProcessorTest.java
@@ -1,45 +1,66 @@
 package software.amazon.lambda.powertools.core.internal;
 
+import com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.RequestHandler;
+import com.amazonaws.services.lambda.runtime.RequestStreamHandler;
 import org.aspectj.lang.ProceedingJoinPoint;
 import org.aspectj.lang.Signature;
 import org.junit.jupiter.api.Test;
 
+import java.io.InputStream;
+import java.io.OutputStream;
+
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
 class LambdaHandlerProcessorTest {
 
     @Test
-    void shouldTreatProfilerHandlerMethodAsValid() {
-        ProceedingJoinPoint pjpMock = mock(ProceedingJoinPoint.class);
-        Signature signature = mock(Signature.class);
-        when(signature.getName()).thenReturn("requestHandler");
-        when(pjpMock.getSignature()).thenReturn(signature);
+    void isHandlerMethod_shouldRecognizeRequestHandler() {
+        ProceedingJoinPoint pjpMock = mockRequestHandlerPjp();
 
-        assertThat(LambdaHandlerProcessor.isHandlerMethod(pjpMock))
-                .isTrue();
+        assertThat(LambdaHandlerProcessor.isHandlerMethod(pjpMock)).isTrue();
     }
 
     @Test
-    void shouldTreatDefaultHandlerMethodAsValid() {
-        ProceedingJoinPoint pjpMock = mock(ProceedingJoinPoint.class);
-        Signature signature = mock(Signature.class);
-        when(signature.getName()).thenReturn("handleRequest");
-        when(pjpMock.getSignature()).thenReturn(signature);
+    void isHandlerMethod_shouldRecognizeRequestStreamHandler() {
+        ProceedingJoinPoint pjpMock = mockRequestStreamHandlerPjp();
 
-        assertThat(LambdaHandlerProcessor.isHandlerMethod(pjpMock))
-                .isTrue();
+        assertThat(LambdaHandlerProcessor.isHandlerMethod(pjpMock)).isTrue();
     }
 
     @Test
-    void shouldNotTreatOtherMethodNamesAsValidHandlerMethod() {
-        ProceedingJoinPoint pjpMock = mock(ProceedingJoinPoint.class);
+    void placedOnRequestHandler_shouldRecognizeRequestHandler() {
+        ProceedingJoinPoint pjpMock = mockRequestHandlerPjp();
+
+        assertThat(LambdaHandlerProcessor.placedOnRequestHandler(pjpMock)).isTrue();
+    }
+
+    @Test
+    void placedOnStreamHandler_shouldRecognizeRequestStreamHandler() {
+        ProceedingJoinPoint pjpMock = mockRequestStreamHandlerPjp();
+
+        assertThat(LambdaHandlerProcessor.placedOnStreamHandler(pjpMock)).isTrue();
+    }
+
+    private static ProceedingJoinPoint mockRequestHandlerPjp() {
         Signature signature = mock(Signature.class);
-        when(signature.getName()).thenReturn("handleRequestInvalid");
+        when(signature.getDeclaringType()).thenReturn(RequestHandler.class);
+        ProceedingJoinPoint pjpMock = mock(ProceedingJoinPoint.class);
+        Object[] args = {new Object(), mock(Context.class)};
+        when(pjpMock.getArgs()).thenReturn(args);
         when(pjpMock.getSignature()).thenReturn(signature);
+        return pjpMock;
+    }
 
-        assertThat(LambdaHandlerProcessor.isHandlerMethod(pjpMock))
-                .isFalse();
+    private static ProceedingJoinPoint mockRequestStreamHandlerPjp() {
+        Signature signature = mock(Signature.class);
+        when(signature.getDeclaringType()).thenReturn(RequestStreamHandler.class);
+        ProceedingJoinPoint pjpMock = mock(ProceedingJoinPoint.class);
+        Object[] args = {mock(InputStream.class), mock(OutputStream.class), mock(Context.class)};
+        when(pjpMock.getArgs()).thenReturn(args);
+        when(pjpMock.getSignature()).thenReturn(signature);
+        return pjpMock;
     }
 }
diff --git a/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/internal/IdempotentAspect.java b/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/internal/IdempotentAspect.java
@@ -30,7 +30,6 @@
 import java.lang.annotation.Annotation;
 import java.lang.reflect.Method;
 
-import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.isHandlerMethod;
 import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.placedOnRequestHandler;
 
 /**
@@ -58,7 +57,7 @@ public Object around(ProceedingJoinPoint pjp,
             throw new IdempotencyConfigurationException("The annotated method doesn't return anything. Unable to perform idempotency on void return type");
         }
 
-        boolean isHandler = (isHandlerMethod(pjp) && placedOnRequestHandler(pjp));
+        boolean isHandler = placedOnRequestHandler(pjp);
         JsonNode payload = getPayload(pjp, method, isHandler);
         if (payload == null) {
             throw new IdempotencyConfigurationException("Unable to get payload from the method. Ensure there is at least one parameter or that you use @IdempotencyKey");
diff --git a/powertools-metrics/src/main/java/software/amazon/lambda/powertools/metrics/internal/LambdaMetricsAspect.java b/powertools-metrics/src/main/java/software/amazon/lambda/powertools/metrics/internal/LambdaMetricsAspect.java
@@ -22,8 +22,6 @@
 import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.extractContext;
 import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.isColdStart;
 import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.isHandlerMethod;
-import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.placedOnRequestHandler;
-import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.placedOnStreamHandler;
 import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.serviceName;
 import static software.amazon.lambda.powertools.metrics.MetricsUtils.hasDefaultDimension;
 import static software.amazon.lambda.powertools.metrics.MetricsUtils.metricsLogger;
@@ -44,9 +42,7 @@ public Object around(ProceedingJoinPoint pjp,
                          Metrics metrics) throws Throwable {
         Object[] proceedArgs = pjp.getArgs();
 
-        if (isHandlerMethod(pjp)
-                && (placedOnRequestHandler(pjp)
-                || placedOnStreamHandler(pjp))) {
+        if (isHandlerMethod(pjp)) {
 
             MetricsLogger logger = metricsLogger();
 
diff --git a/powertools-tracing/src/main/java/software/amazon/lambda/powertools/tracing/internal/LambdaTracingAspect.java b/powertools-tracing/src/main/java/software/amazon/lambda/powertools/tracing/internal/LambdaTracingAspect.java
@@ -26,8 +26,6 @@
 import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.isColdStart;
 import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.isHandlerMethod;
 import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.isSamLocal;
-import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.placedOnRequestHandler;
-import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.placedOnStreamHandler;
 import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.serviceName;
 import static software.amazon.lambda.powertools.tracing.TracingUtils.objectMapper;
 
@@ -48,7 +46,7 @@ public Object around(ProceedingJoinPoint pjp,
                 () -> "## " + pjp.getSignature().getName()));
         segment.setNamespace(namespace(tracing));
 
-        if (placedOnHandlerMethod(pjp)) {
+        if (isHandlerMethod(pjp)) {
             segment.putAnnotation("ColdStart", isColdStart());
             segment.putAnnotation("Service", namespace(tracing));
         }
@@ -62,7 +60,7 @@ public Object around(ProceedingJoinPoint pjp,
                 segment.putMetadata(namespace(tracing), pjp.getSignature().getName() + " response", null != objectMapper() ? objectMapper().writeValueAsString(methodReturn): methodReturn);
             }
 
-            if (placedOnHandlerMethod(pjp)) {
+            if (isHandlerMethod(pjp)) {
                 coldStartDone();
             }
 
@@ -116,11 +114,6 @@ private String namespace(Tracing powerToolsTracing) {
         return powerToolsTracing.namespace().isEmpty() ? serviceName() : powerToolsTracing.namespace();
     }
 
-    private boolean placedOnHandlerMethod(ProceedingJoinPoint pjp) {
-        return isHandlerMethod(pjp)
-                && (placedOnRequestHandler(pjp) || placedOnStreamHandler(pjp));
-    }
-
     private boolean environmentVariable(String key) {
         return Boolean.parseBoolean(SystemWrapper.getenv(key));
     }
diff --git a/powertools-validation/src/main/java/software/amazon/lambda/powertools/validation/internal/ValidationAspect.java b/powertools-validation/src/main/java/software/amazon/lambda/powertools/validation/internal/ValidationAspect.java
@@ -26,7 +26,6 @@
 
 import static com.networknt.schema.SpecVersion.VersionFlag.V201909;
 import static java.nio.charset.StandardCharsets.UTF_8;
-import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.isHandlerMethod;
 import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.placedOnRequestHandler;
 import static software.amazon.lambda.powertools.utilities.jmespath.Base64Function.decode;
 import static software.amazon.lambda.powertools.utilities.jmespath.Base64GZipFunction.decompress;
@@ -55,8 +54,7 @@ public Object around(ProceedingJoinPoint pjp,
             ValidationConfig.get().setSchemaVersion(validation.schemaVersion());
         }
 
-        if (isHandlerMethod(pjp)
-                && placedOnRequestHandler(pjp)) {
+        if (placedOnRequestHandler(pjp)) {
             validationNeeded = true;
 
             if (!validation.inboundSchema().isEmpty()) {

Original file line number	Diff line number	Diff line change
`@@ -37,9 +37,7 @@ private LambdaHandlerProcessor() {`
`37`	`37`	`}`
`38`	`38`
`39`	`39`	`public static boolean isHandlerMethod(final ProceedingJoinPoint pjp) {`
`40`		`- return "handleRequest".equals(pjp.getSignature().getName()) \|\|`
`41`		`- // https://docs.aws.amazon.com/codeguru/latest/profiler-ug/lambda-custom.html`
`42`		`- "requestHandler".equals(pjp.getSignature().getName());`
	`40`	`+ return placedOnRequestHandler(pjp) \|\| placedOnStreamHandler(pjp);`
`43`	`41`	`}`
`44`	`42`
`45`	`43`	`public static boolean placedOnRequestHandler(final ProceedingJoinPoint pjp) {`