add utility methods for callback

Author: zhongkechen

examples/src/test/java/software/amazon/lambda/durable/examples/CloudBasedIntegrationTest.java

@@ -15,9 +15,9 @@
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.CsvSource;
 import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider;
-import software.amazon.awssdk.core.SdkBytes;
 import software.amazon.awssdk.regions.Region;
 import software.amazon.awssdk.services.lambda.LambdaClient;
+import software.amazon.awssdk.services.lambda.model.ErrorObject;
 import software.amazon.awssdk.services.lambda.model.OperationStatus;
 import software.amazon.awssdk.services.sts.StsClient;
 import software.amazon.lambda.durable.TypeToken;
@@ -344,8 +344,7 @@ void testCallbackExample() {
         // Complete the preapproval callback
         execution.pollUntil(exec -> exec.hasCallback("preapproval-callback"));
         var preapprovalCallbackId = execution.getCallbackId("preapproval-callback");
-        lambdaClient.sendDurableExecutionCallbackSuccess(
-                req -> req.callbackId(preapprovalCallbackId).result(SdkBytes.fromUtf8String("\"preapproved\"")));
+        execution.completeCallback(preapprovalCallbackId, "\"preapproved\"");
 
         // Wait for callback to appear
         execution.pollUntil(exec -> exec.hasCallback("approval"));
@@ -355,8 +354,7 @@ void testCallbackExample() {
         assertNotNull(callbackId);
 
         // Complete the callback using AWS SDK
-        lambdaClient.sendDurableExecutionCallbackSuccess(
-                req -> req.callbackId(callbackId).result(SdkBytes.fromUtf8String("\"approved\"")));
+        execution.completeCallback(callbackId, "\"approved\"");
 
         // Wait for execution to complete
         var result = execution.pollUntilComplete();
@@ -385,8 +383,7 @@ void testCallbackExampleWithFailure() {
 
         execution.pollUntil(exec -> exec.hasCallback("preapproval-callback"));
         var preapprovalCallbackId = execution.getCallbackId("preapproval-callback");
-        lambdaClient.sendDurableExecutionCallbackSuccess(
-                req -> req.callbackId(preapprovalCallbackId).result(SdkBytes.fromUtf8String("\"preapproved\"")));
+        execution.completeCallback(preapprovalCallbackId, "\"preapproved\"");
 
         // Wait for callback to appear
         execution.pollUntil(exec -> exec.hasCallback("approval"));
@@ -396,8 +393,12 @@ void testCallbackExampleWithFailure() {
         assertNotNull(callbackId);
 
         // Fail the callback using AWS SDK
-        lambdaClient.sendDurableExecutionCallbackFailure(req -> req.callbackId(callbackId)
-                .error(err -> err.errorType("ApprovalRejected").errorMessage("Approval rejected by manager")));
+        execution.failCallback(
+                callbackId,
+                ErrorObject.builder()
+                        .errorType("ApprovalRejected")
+                        .errorMessage("Approval rejected by manager")
+                        .build());
 
         // Wait for execution to complete
         var result = execution.pollUntilComplete();
@@ -423,8 +424,7 @@ void testCallbackExampleWithTimeout() {
 
         execution.pollUntil(exec -> exec.hasCallback("preapproval-callback"));
         var preapprovalCallbackId = execution.getCallbackId("preapproval-callback");
-        lambdaClient.sendDurableExecutionCallbackSuccess(
-                req -> req.callbackId(preapprovalCallbackId).result(SdkBytes.fromUtf8String("\"preapproved\"")));
+        execution.completeCallback(preapprovalCallbackId, "\"preapproved\"");
 
         // Wait for callback to appear
         execution.pollUntil(exec -> exec.hasCallback("approval"));
@@ -454,8 +454,9 @@ void testCallbackExampleWithWaitForCallbackFailure() {
 
         execution.pollUntil(exec -> exec.hasCallback("preapproval-callback"));
         var preapprovalCallbackId = execution.getCallbackId("preapproval-callback");
-        lambdaClient.sendDurableExecutionCallbackFailure(
-                req -> req.callbackId(preapprovalCallbackId).error(err -> err.errorMessage("preapproval denied")));
+        execution.failCallback(
+                preapprovalCallbackId,
+                ErrorObject.builder().errorMessage("preapproval denied").build());
 
         // Wait for callback to appear
         execution.pollUntil(exec -> exec.hasCallback("approval"));

sdk-integration-tests/src/test/java/software/amazon/lambda/durable/DurableExecutionCheckpointTest.java

@@ -10,7 +10,7 @@
 import software.amazon.lambda.durable.execution.DurableExecutor;
 import software.amazon.lambda.durable.model.DurableExecutionInput;
 import software.amazon.lambda.durable.model.ExecutionStatus;
-import software.amazon.lambda.durable.testing.LocalMemoryExecutionClient;
+import software.amazon.lambda.durable.testing.local.LocalMemoryExecutionClient;
 
 /** Integration tests that verify checkpoint behavior using LocalMemoryExecutionClient */
 class DurableExecutionCheckpointTest {

sdk-testing/src/main/java/software/amazon/lambda/durable/testing/AsyncExecution.java

@@ -6,13 +6,16 @@
 import java.time.Instant;
 import java.util.List;
 import java.util.function.Predicate;
+import software.amazon.awssdk.core.SdkBytes;
 import software.amazon.awssdk.services.lambda.LambdaClient;
+import software.amazon.awssdk.services.lambda.model.ErrorObject;
 import software.amazon.awssdk.services.lambda.model.Event;
 import software.amazon.awssdk.services.lambda.model.EventType;
 import software.amazon.awssdk.services.lambda.model.GetDurableExecutionHistoryRequest;
 import software.amazon.awssdk.services.lambda.model.ResourceNotFoundException;
 import software.amazon.lambda.durable.TypeToken;
 import software.amazon.lambda.durable.model.ExecutionStatus;
+import software.amazon.lambda.durable.testing.cloud.HistoryEventProcessor;
 
 /**
  * Handle for an asynchronously executing durable function. Allows incremental polling and inspection of execution
@@ -163,6 +166,23 @@ public String getExecutionArn() {
         return executionArn;
     }
 
+    /** calls sendDurableExecutionCallbackSuccess with the given callbackId and result */
+    public void completeCallback(String callbackId, String result) {
+        lambdaClient.sendDurableExecutionCallbackSuccess(
+                req -> req.callbackId(callbackId).result(SdkBytes.fromUtf8String(result)));
+    }
+
+    /** calls sendDurableExecutionCallbackFailure with the give callbackId and error */
+    public void failCallback(String callbackId, ErrorObject error) {
+        lambdaClient.sendDurableExecutionCallbackFailure(
+                req -> req.callbackId(callbackId).error(error));
+    }
+
+    /** call sendDurableExecutionCallbackHeartbeat API with the give callbackId */
+    public void heartbeatCallback(String callbackId) {
+        lambdaClient.sendDurableExecutionCallbackHeartbeat(req -> req.callbackId(callbackId));
+    }
+
     private void refreshHistory() {
         try {
             var request = GetDurableExecutionHistoryRequest.builder()

sdk-testing/src/main/java/software/amazon/lambda/durable/testing/CloudDurableTestRunner.java

@@ -10,6 +10,8 @@
 import software.amazon.awssdk.services.lambda.model.InvokeRequest;
 import software.amazon.lambda.durable.TypeToken;
 import software.amazon.lambda.durable.serde.JacksonSerDes;
+import software.amazon.lambda.durable.testing.cloud.HistoryEventProcessor;
+import software.amazon.lambda.durable.testing.cloud.HistoryPoller;
 
 /**
  * Test runner for durable Lambda functions deployed to AWS. Invokes a real Lambda function, polls execution history,

sdk-testing/src/main/java/software/amazon/lambda/durable/testing/LocalDurableTestRunner.java

@@ -22,6 +22,8 @@
 import software.amazon.lambda.durable.model.ExecutionStatus;
 import software.amazon.lambda.durable.serde.JacksonSerDes;
 import software.amazon.lambda.durable.serde.SerDes;
+import software.amazon.lambda.durable.testing.local.LocalMemoryExecutionClient;
+import software.amazon.lambda.durable.testing.local.OperationResult;
 
 /**
  * In-memory test runner for durable Lambda functions. Simulates the Lambda re-invocation loop locally without requiring

sdk-testing/src/main/java/software/amazon/lambda/durable/testing/TestOperation.java

@@ -26,11 +26,11 @@ public class TestOperation {
     private final List<Event> events;
     private final SerDes serDes;
 
-    TestOperation(Operation operation, SerDes serDes) {
+    public TestOperation(Operation operation, SerDes serDes) {
         this(operation, List.of(), serDes);
     }
 
-    TestOperation(Operation operation, List<Event> events, SerDes serDes) {
+    public TestOperation(Operation operation, List<Event> events, SerDes serDes) {
         this.operation = operation;
         this.events = events;
         this.serDes = serDes;

sdk-testing/src/main/java/software/amazon/lambda/durable/testing/TestResult.java

@@ -32,7 +32,7 @@ public class TestResult<O> {
     private final List<Event> allEvents;
     private final SerDes serDes;
 
-    TestResult(
+    public TestResult(
             ExecutionStatus status,
             String resultPayload,
             ErrorObject error,

…rable/testing/HistoryEventProcessor.java …testing/cloud/HistoryEventProcessor.javasdk-testing/src/main/java/software/amazon/lambda/durable/testing/HistoryEventProcessor.java renamed to sdk-testing/src/main/java/software/amazon/lambda/durable/testing/cloud/HistoryEventProcessor.java sdk-testing/src/main/java/software/amazon/lambda/durable/testing/HistoryEventProcessor.java renamed to sdk-testing/src/main/java/software/amazon/lambda/durable/testing/cloud/HistoryEventProcessor.java

@@ -1,6 +1,6 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
-package software.amazon.lambda.durable.testing;
+package software.amazon.lambda.durable.testing.cloud;
 
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -18,6 +18,10 @@
 import software.amazon.lambda.durable.TypeToken;
 import software.amazon.lambda.durable.model.ExecutionStatus;
 import software.amazon.lambda.durable.serde.JacksonSerDes;
+import software.amazon.lambda.durable.testing.AsyncExecution;
+import software.amazon.lambda.durable.testing.CloudDurableTestRunner;
+import software.amazon.lambda.durable.testing.TestOperation;
+import software.amazon.lambda.durable.testing.TestResult;
 
 /**
  * Processes execution history events from the GetDurableExecutionHistory API into {@link TestResult} objects. Used by

…ambda/durable/testing/HistoryPoller.java …durable/testing/cloud/HistoryPoller.javasdk-testing/src/main/java/software/amazon/lambda/durable/testing/HistoryPoller.java renamed to sdk-testing/src/main/java/software/amazon/lambda/durable/testing/cloud/HistoryPoller.java sdk-testing/src/main/java/software/amazon/lambda/durable/testing/HistoryPoller.java renamed to sdk-testing/src/main/java/software/amazon/lambda/durable/testing/cloud/HistoryPoller.java

@@ -1,6 +1,6 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
-package software.amazon.lambda.durable.testing;
+package software.amazon.lambda.durable.testing.cloud;
 
 import java.time.Duration;
 import java.time.Instant;
@@ -10,6 +10,7 @@
 import software.amazon.awssdk.services.lambda.model.Event;
 import software.amazon.awssdk.services.lambda.model.EventType;
 import software.amazon.awssdk.services.lambda.model.GetDurableExecutionHistoryRequest;
+import software.amazon.lambda.durable.testing.CloudDurableTestRunner;
 
 /**
  * Polls the GetDurableExecutionHistory API until execution completes or a timeout is reached. Used by

…mbda/durable/testing/EventProcessor.java …urable/testing/local/EventProcessor.javasdk-testing/src/main/java/software/amazon/lambda/durable/testing/EventProcessor.java renamed to sdk-testing/src/main/java/software/amazon/lambda/durable/testing/local/EventProcessor.java sdk-testing/src/main/java/software/amazon/lambda/durable/testing/EventProcessor.java renamed to sdk-testing/src/main/java/software/amazon/lambda/durable/testing/local/EventProcessor.java

@@ -1,6 +1,6 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
-package software.amazon.lambda.durable.testing;
+package software.amazon.lambda.durable.testing.local;
 
 import static software.amazon.awssdk.services.lambda.model.EventType.*;