Support handler methods without any argument (#276)

--------- Co-authored-by: Andre Onuki <[email protected]>
newrelic · Dec 24, 2024 · edccdf3 · edccdf3
1 parent 02f7d0f
commit edccdf3
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 11 deletions.
diff --git a/java/src/main/java/com/newrelic/java/JavaClassLoader.java b/java/src/main/java/com/newrelic/java/JavaClassLoader.java
@@ -33,13 +33,17 @@ private interface UnsafeHandler {
 
     static JavaClassLoader initializeRequestHandler(Class<?> loadedClass, String methodName) throws ReflectiveOperationException {
         Class<?> methodReturnType = Object.class;
-        Class<?> methodInputType = Object.class;
+        Class<?> methodInputType = null;
         Class<?> methodContextType = null;
+        int numberOfArguments = 0;
         for (Method method : loadedClass.getMethods()) {
             if (isUserHandlerMethod(method, methodName, loadedClass)) {
                 methodReturnType = method.getReturnType();
-                methodInputType = method.getParameterTypes()[0];
-                if (method.getParameterTypes().length == 2) {
+                numberOfArguments = method.getParameterTypes().length;
+                if (numberOfArguments == 1) {
+                    methodInputType = method.getParameterTypes()[0];
+                } else if (numberOfArguments == 2) {
+                    methodInputType = method.getParameterTypes()[0];
                     methodContextType = method.getParameterTypes()[1];
                 }
                 break;
@@ -54,22 +58,28 @@ static JavaClassLoader initializeRequestHandler(Class<?> loadedClass, String met
                 getMethodType(methodReturnType, methodInputType, methodContextType)
         ).bindTo(classInstance);
 
-        return new JavaClassLoader(methodInputType, methodHandle, methodContextType != null);
+        return new JavaClassLoader(methodInputType, methodHandle, numberOfArguments);
     }
 
     private static MethodType getMethodType(Class<?> methodReturnType, Class<?> methodInputType, Class<?> methodContextType) {
-        if (methodContextType == null) {
+        if (methodInputType == null) {
+            return MethodType.methodType(methodReturnType);
+        } else if (methodContextType == null) {
             return MethodType.methodType(methodReturnType, methodInputType);
         }
         return MethodType.methodType(methodReturnType, methodInputType, methodContextType);
     }
 
     // RequestHandler implementation constructor
-    private JavaClassLoader(Class<?> inputType, MethodHandle methodHandle, boolean hasTwoArguments) {
+    private JavaClassLoader(Class<?> inputType, MethodHandle methodHandle, int numberOfArguments) {
         this.inputType = inputType;
-        this.executor = hasTwoArguments
-                ? methodHandle::invokeWithArguments
-                : (handlerType, contextParam) -> methodHandle.invokeWithArguments(handlerType);
+        if (numberOfArguments == 0) {
+            this.executor = (input, context) -> methodHandle.invoke();
+        } else if (numberOfArguments == 1) {
+            this.executor = (input, context) -> methodHandle.invokeWithArguments(input);
+        } else {
+            this.executor = methodHandle::invokeWithArguments;
+        }
     }
 
     @Override
@@ -94,7 +104,7 @@ private static boolean isUserHandlerMethod(Method method, String methodName, Cla
             return false;
         }
 
-        if (method.getParameterTypes().length == 1) {
+        if (method.getParameterTypes().length <= 1) {
             return true;
         }
 
@@ -104,7 +114,7 @@ private static boolean isUserHandlerMethod(Method method, String methodName, Cla
     }
 
     private Object mappingInputToHandlerType(Object inputParam, Class<?> inputType) throws JsonProcessingException {
-        if (inputType.isAssignableFrom(Number.class) || inputType.isAssignableFrom(String.class)) {
+        if (inputType == null || inputType.isAssignableFrom(Number.class) || inputType.isAssignableFrom(String.class)) {
             return inputParam;
         } else if (LambdaEventSerializers.isLambdaSupportedEvent(inputType.getName())) {
             PojoSerializer<?> serializer = LambdaEventSerializers.serializerFor(inputType, JavaClassLoader.class.getClassLoader());

diff --git a/java/src/test/java/com/newrelic/java/JavaClassLoaderTest.java b/java/src/test/java/com/newrelic/java/JavaClassLoaderTest.java
@@ -51,4 +51,10 @@ public void testWithStringRequestHandler() throws ReflectiveOperationException {
         Assert.assertEquals("Hello World", loader.handleRequest(STRING_INPUT, null));
     }
 
+    @Test
+    public void testWithoutArgumentRequestHandler() throws ReflectiveOperationException {
+        JavaClassLoader loader = JavaClassLoader.initializeRequestHandler(PojoHandlerWithNoArgument.class, "handleRequest");
+        Assert.assertEquals("Hello World!", loader.handleRequest(INPUT_AS_OBJECT, null));
+    }
+
 }
diff --git a/java/src/test/java/com/newrelic/java/PojoHandlerWithNoArgument.java b/java/src/test/java/com/newrelic/java/PojoHandlerWithNoArgument.java
@@ -0,0 +1,10 @@
+package com.newrelic.java;
+
+public class PojoHandlerWithNoArgument {
+
+
+    public String handleRequest() {
+        return "Hello World!";
+    }
+
+}