Fix bug: We need to pass the return type explicitly to support deserialization to the return type of the method.

Author: phipag

powertools-idempotency/powertools-idempotency-core/src/main/java/software/amazon/lambda/powertools/idempotency/PowertoolsIdempotency.java

@@ -30,39 +30,48 @@
  * <p>This class is thread-safe. All operations delegate to the underlying persistence store
  * which handles concurrent access safely.</p>
  * 
- * <p>Example usage with Supplier:</p>
+ * <p><strong>Important:</strong> Always call {@link Idempotency#registerLambdaContext(Context)}
+ * at the start of your handler to enable proper timeout handling.</p>
+ * 
+ * <p>Example usage with Function (single parameter):</p>
  * <pre>{@code
- * public List<String> handleRequest(SQSEvent event, Context context) {
+ * public Basket handleRequest(Product input, Context context) {
  *     Idempotency.registerLambdaContext(context);
- *     return event.getRecords().stream()
- *         .map(record -> PowertoolsIdempotency.makeIdempotent(
- *             record.getBody(), // used as idempotency key
- *             () -> processPayment(record.getMessageId(), record.getBody())))
- *         .collect(Collectors.toList());
+ *     return PowertoolsIdempotency.makeIdempotent(this::processProduct, input, Basket.class);
+ * }
+ * 
+ * private Basket processProduct(Product product) {
+ *     // business logic
  * }
  * }</pre>
  * 
- * <p>Example usage with Function:</p>
+ * <p>Example usage with Supplier (multi-parameter methods):</p>
  * <pre>{@code
- * public Basket handleRequest(Product input, Context context) {
+ * public String handleRequest(SQSEvent event, Context context) {
  *     Idempotency.registerLambdaContext(context);
- *     return PowertoolsIdempotency.makeIdempotent(this::process, input);
+ *     return PowertoolsIdempotency.makeIdempotent(
+ *         event.getRecords().get(0).getBody(),
+ *         () -> processPayment(orderId, amount, currency),
+ *         String.class
+ *     );
  * }
  * }</pre>
  * 
- * <p>When different methods use the same payload as idempotency key, use
- * {@link #makeIdempotent(String, Object, Supplier)} with explicit function names
+ * <p>When different methods use the same payload as idempotency key, use explicit function names
  * to differentiate between them:</p>
  * <pre>{@code
  * // Different methods, same payload
- * PowertoolsIdempotency.makeIdempotent("processPayment", orderId, () -> processPayment(orderId));
- * PowertoolsIdempotency.makeIdempotent("refundPayment", orderId, () -> refundPayment(orderId));
+ * PowertoolsIdempotency.makeIdempotent("processPayment", orderId, 
+ *     () -> processPayment(orderId), String.class);
+ * 
+ * PowertoolsIdempotency.makeIdempotent("refundPayment", orderId, 
+ *     () -> refundPayment(orderId), String.class);
  * }</pre>
  * 
  * @see Idempotency
- * @see #makeIdempotent(Object, Supplier)
- * @see #makeIdempotent(String, Object, Supplier)
- * @see #makeIdempotent(Function, Object)
+ * @see #makeIdempotent(Object, Supplier, Class)
+ * @see #makeIdempotent(String, Object, Supplier, Class)
+ * @see #makeIdempotent(Function, Object, Class)
  */
 public final class PowertoolsIdempotency {
 
@@ -80,20 +89,22 @@ private PowertoolsIdempotency() {
      * such as batch processors.</p>
      * 
      * <p>This method is suitable for making methods idempotent that have more than one parameter.
-     * For simple single-parameter methods, {@link #makeIdempotent(Function, Object)} is more intuitive.</p>
+     * For simple single-parameter methods, {@link #makeIdempotent(Function, Object, Class)} is more intuitive.</p>
      * 
      * <p><strong>Note:</strong> If you need to call different functions with the same payload,
-     * use {@link #makeIdempotent(String, Object, Supplier)} to specify distinct function names.
+     * use {@link #makeIdempotent(String, Object, Supplier, Class)} to specify distinct function names.
      * This ensures each function has its own idempotency scope.</p>
      * 
      * @param idempotencyKey the key used for idempotency (will be converted to JSON)
      * @param function the function to make idempotent
+     * @param returnType the class of the return type for deserialization
      * @param <T> the return type of the function
      * @return the result of the function execution (either fresh or cached)
      * @throws Throwable if the function execution fails
      */
-    public static <T> T makeIdempotent(Object idempotencyKey, Supplier<T> function) throws Throwable {
-        return makeIdempotent(DEFAULT_FUNCTION_NAME, idempotencyKey, function);
+    public static <T> T makeIdempotent(Object idempotencyKey, Supplier<T> function, Class<T> returnType)
+            throws Throwable {
+        return makeIdempotent(DEFAULT_FUNCTION_NAME, idempotencyKey, function, returnType);
     }
 
     /**
@@ -102,25 +113,24 @@ public static <T> T makeIdempotent(Object idempotencyKey, Supplier<T> function)
      * <p>This method is thread-safe and can be used in parallel processing scenarios
      * such as batch processors.</p>
      * 
-     * <p>Note: The return type is inferred from the actual result. For cached responses,
-     * deserialization uses the runtime type of the stored result.</p>
-     * 
      * @param functionName the name of the function (used for persistence store configuration)
      * @param idempotencyKey the key used for idempotency (will be converted to JSON)
      * @param function the function to make idempotent
+     * @param returnType the class of the return type for deserialization
      * @param <T> the return type of the function
      * @return the result of the function execution (either fresh or cached)
      * @throws Throwable if the function execution fails
      */
     @SuppressWarnings("unchecked")
-    public static <T> T makeIdempotent(String functionName, Object idempotencyKey, Supplier<T> function)
+    public static <T> T makeIdempotent(String functionName, Object idempotencyKey, Supplier<T> function,
+            Class<T> returnType)
             throws Throwable {
         JsonNode payload = JsonConfig.get().getObjectMapper().valueToTree(idempotencyKey);
         Context lambdaContext = Idempotency.getInstance().getConfig().getLambdaContext();
 
         IdempotencyHandler handler = new IdempotencyHandler(
                 function::get,
-                Object.class,
+                returnType,
                 functionName,
                 payload,
                 lambdaContext);
@@ -133,20 +143,21 @@ public static <T> T makeIdempotent(String functionName, Object idempotencyKey, S
      * Makes a function with one parameter idempotent.
      * The parameter is used as the idempotency key.
      * 
-     * <p>For functions with more than one parameter, use {@link #makeIdempotent(Object, Supplier)} instead.</p>
+     * <p>For functions with more than one parameter, use {@link #makeIdempotent(Object, Supplier, Class)} instead.</p>
      * 
      * <p><strong>Note:</strong> If you need to call different functions with the same argument,
-     * use {@link #makeIdempotent(String, Object, Supplier)} to specify distinct function names.</p>
+     * use {@link #makeIdempotent(String, Object, Supplier, Class)} to specify distinct function names.</p>
      * 
      * @param function the function to make idempotent (method reference)
      * @param arg the argument to pass to the function (also used as idempotency key)
+     * @param returnType the class of the return type for deserialization
      * @param <T> the argument type
      * @param <R> the return type
      * @return the result of the function execution (either fresh or cached)
      * @throws Throwable if the function execution fails
      */
-    public static <T, R> R makeIdempotent(Function<T, R> function, T arg) throws Throwable {
-        return makeIdempotent(DEFAULT_FUNCTION_NAME, arg, () -> function.apply(arg));
+    public static <T, R> R makeIdempotent(Function<T, R> function, T arg, Class<R> returnType) throws Throwable {
+        return makeIdempotent(DEFAULT_FUNCTION_NAME, arg, () -> function.apply(arg), returnType);
     }
 
 }

powertools-idempotency/powertools-idempotency-core/src/test/java/software/amazon/lambda/powertools/idempotency/PowertoolsIdempotencyTest.java

@@ -19,6 +19,9 @@
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 
+import java.time.Instant;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.OptionalInt;
 
 import org.junit.jupiter.api.Test;
@@ -31,11 +34,14 @@
 import com.fasterxml.jackson.databind.JsonNode;
 
 import software.amazon.lambda.powertools.common.stubs.TestLambdaContext;
+import software.amazon.lambda.powertools.idempotency.exceptions.IdempotencyItemAlreadyExistsException;
+import software.amazon.lambda.powertools.idempotency.exceptions.IdempotencyItemNotFoundException;
 import software.amazon.lambda.powertools.idempotency.handlers.PowertoolsIdempotencyFunction;
 import software.amazon.lambda.powertools.idempotency.handlers.PowertoolsIdempotencyMultiArgFunction;
 import software.amazon.lambda.powertools.idempotency.model.Basket;
 import software.amazon.lambda.powertools.idempotency.model.Product;
 import software.amazon.lambda.powertools.idempotency.persistence.BasePersistenceStore;
+import software.amazon.lambda.powertools.idempotency.persistence.DataRecord;
 
 @ExtendWith(MockitoExtension.class)
 class PowertoolsIdempotencyTest {
@@ -83,7 +89,8 @@ void testMakeIdempotentWithFunctionName() throws Throwable {
                 .configure();
         Idempotency.registerLambdaContext(context);
 
-        String result = PowertoolsIdempotency.makeIdempotent("myFunction", "test-key", () -> "test-result");
+        String result = PowertoolsIdempotency.makeIdempotent("myFunction", "test-key", () -> "test-result",
+                String.class);
 
         assertThat(result).isEqualTo("test-result");
 
@@ -100,7 +107,7 @@ void testMakeIdempotentWithMethodReferenceUsesDefaultName() throws Throwable {
                 .configure();
         Idempotency.registerLambdaContext(context);
 
-        String result = PowertoolsIdempotency.makeIdempotent("test-key", this::helperMethod);
+        String result = PowertoolsIdempotency.makeIdempotent("test-key", this::helperMethod, String.class);
 
         assertThat(result).isEqualTo("helper-result");
 
@@ -122,7 +129,7 @@ void testMakeIdempotentWithFunctionOverload() throws Throwable {
         Idempotency.registerLambdaContext(context);
 
         Product p = new Product(42, "test product", 10);
-        Basket result = PowertoolsIdempotency.makeIdempotent(this::processProduct, p);
+        Basket result = PowertoolsIdempotency.makeIdempotent(this::processProduct, p, Basket.class);
 
         assertThat(result.getProducts()).hasSize(1);
         assertThat(result.getProducts().get(0)).isEqualTo(p);
@@ -166,6 +173,66 @@ void firstCall_withExplicitIdempotencyKey_shouldPutInStore() {
         assertThat(resultCaptor.getValue()).isEqualTo(basket);
     }
 
+    @Test
+    void secondCall_shouldRetrieveFromCacheAndDeserialize() throws Throwable {
+        // Use in-memory persistence store to test actual serialization/deserialization
+        Map<String, DataRecord> data = new HashMap<>();
+        BasePersistenceStore inMemoryStore = new BasePersistenceStore() {
+            @Override
+            public DataRecord getRecord(String idempotencyKey) throws IdempotencyItemNotFoundException {
+                DataRecord dr = data.get(idempotencyKey);
+                if (dr == null) {
+                    throw new IdempotencyItemNotFoundException(idempotencyKey);
+                }
+                return dr;
+            }
+
+            @Override
+            public void putRecord(DataRecord dr, Instant now) throws IdempotencyItemAlreadyExistsException {
+                if (data.containsKey(dr.getIdempotencyKey())) {
+                    throw new IdempotencyItemAlreadyExistsException();
+                }
+                data.put(dr.getIdempotencyKey(), dr);
+            }
+
+            @Override
+            public void updateRecord(DataRecord dr) {
+                data.put(dr.getIdempotencyKey(), dr);
+            }
+
+            @Override
+            public void deleteRecord(String idempotencyKey) {
+                data.remove(idempotencyKey);
+            }
+        };
+
+        Idempotency.config()
+                .withPersistenceStore(inMemoryStore)
+                .configure();
+        Idempotency.registerLambdaContext(context);
+
+        Product p = new Product(42, "test product", 10);
+        int[] callCount = { 0 };
+
+        // First call - executes function and stores result
+        Basket result1 = PowertoolsIdempotency.makeIdempotent(p, () -> {
+            callCount[0]++;
+            return processProduct(p);
+        }, Basket.class);
+        assertThat(result1.getProducts()).hasSize(1);
+        assertThat(callCount[0]).isEqualTo(1);
+
+        // Second call - should retrieve from cache, deserialize, and NOT execute function
+        Basket result2 = PowertoolsIdempotency.makeIdempotent(p, () -> {
+            callCount[0]++;
+            return processProduct(p);
+        }, Basket.class);
+        assertThat(result2.getProducts()).hasSize(1);
+        assertThat(result2.getProducts().get(0).getId()).isEqualTo(42);
+        assertThat(result2.getProducts().get(0).getName()).isEqualTo("test product");
+        assertThat(callCount[0]).isEqualTo(1); // Function should NOT be called again
+    }
+
     @Test
     void concurrentInvocations_shouldNotLeakContext() throws Exception {
         Idempotency.config()

powertools-idempotency/powertools-idempotency-core/src/test/java/software/amazon/lambda/powertools/idempotency/handlers/PowertoolsIdempotencyFunction.java

@@ -37,7 +37,7 @@ public boolean processCalled() {
     public Basket handleRequest(Product input, Context context) {
         Idempotency.registerLambdaContext(context);
         try {
-            return PowertoolsIdempotency.makeIdempotent(this::process, input);
+            return PowertoolsIdempotency.makeIdempotent(this::process, input, Basket.class);
         } catch (Throwable e) {
             throw new RuntimeException(e);
         }

powertools-idempotency/powertools-idempotency-core/src/test/java/software/amazon/lambda/powertools/idempotency/handlers/PowertoolsIdempotencyMultiArgFunction.java

@@ -42,7 +42,8 @@ public String getExtraData() {
     public Basket handleRequest(Product input, Context context) {
         Idempotency.registerLambdaContext(context);
         try {
-            return PowertoolsIdempotency.makeIdempotent(input.getId(), () -> process(input, "extra-data"));
+            return PowertoolsIdempotency.makeIdempotent(input.getId(), () -> process(input, "extra-data"),
+                    Basket.class);
         } catch (Throwable e) {
             throw new RuntimeException(e);
         }