add tests for virtual context

Author: zhongkechen

examples/src/main/java/software/amazon/lambda/durable/examples/child/VirtualChildContextExample.java

@@ -0,0 +1,89 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable.examples.child;
+
+import java.time.Duration;
+import software.amazon.lambda.durable.DurableContext;
+import software.amazon.lambda.durable.DurableFuture;
+import software.amazon.lambda.durable.DurableHandler;
+import software.amazon.lambda.durable.config.RunInChildContextConfig;
+import software.amazon.lambda.durable.examples.types.GreetingRequest;
+
+/**
+ * Example demonstrating virtual child context workflows with the Durable Execution SDK.
+ *
+ * <p>This handler runs three concurrent child contexts using {@code runInChildContextAsync}:
+ *
+ * <ol>
+ *   <li><b>Order validation</b> — performs a step then suspends via {@code wait()} before completing
+ *   <li><b>Inventory check</b> — performs a step then suspends via {@code wait()} before completing
+ *   <li><b>Shipping estimate</b> — nests another child context inside it to demonstrate hierarchical contexts
+ * </ol>
+ *
+ * <p>All three child contexts run concurrently. Results are collected with {@link DurableFuture#allOf} and combined
+ * into a summary string.
+ */
+public class VirtualChildContextExample extends DurableHandler<GreetingRequest, String> {
+
+    @Override
+    public String handleRequest(GreetingRequest input, DurableContext context) {
+        var name = input.getName();
+        context.getLogger().info("Starting child context workflow for {}", name);
+
+        // Child context 1: Order validation — step + wait + step
+        var orderFuture = context.runInChildContextAsync(
+                "order-validation",
+                String.class,
+                child -> {
+                    var prepared = child.step("prepare-order", String.class, stepCtx -> "Order for " + name);
+                    child.getLogger().info("Order prepared, waiting for validation");
+
+                    child.wait("validation-delay", Duration.ofSeconds(5));
+
+                    return child.step("validate-order", String.class, stepCtx -> prepared + " [validated]");
+                },
+                RunInChildContextConfig.builder().isVirtual(true).build());
+
+        // Child context 2: Inventory check — step + wait + step
+        var inventoryFuture = context.runInChildContextAsync(
+                "inventory-check",
+                String.class,
+                child -> {
+                    var stock = child.step("check-stock", String.class, stepCtx -> "Stock available for " + name);
+                    child.getLogger().info("Stock checked, waiting for confirmation");
+
+                    child.wait("confirmation-delay", Duration.ofSeconds(3));
+
+                    return child.step("confirm-inventory", String.class, stepCtx -> stock + " [confirmed]");
+                },
+                RunInChildContextConfig.builder().isVirtual(true).build());
+
+        // Child context 3: Shipping estimate — nests a child context inside it
+        var shippingFuture = context.runInChildContextAsync(
+                "shipping-estimate",
+                String.class,
+                child -> {
+                    var baseRate = child.step("calculate-base-rate", String.class, stepCtx -> "Base rate for " + name);
+
+                    // Nested child context: calculate regional adjustment
+                    var adjustment = child.runInChildContext(
+                            "regional-adjustment",
+                            String.class,
+                            nested -> nested.step(
+                                    "lookup-region", String.class, stepCtx -> baseRate + " + regional adjustment"));
+
+                    return child.step("finalize-shipping", String.class, stepCtx -> adjustment + " [shipping ready]");
+                },
+                RunInChildContextConfig.builder().isVirtual(true).build());
+
+        // Collect all results using allOf
+        context.getLogger().info("Waiting for all child contexts to complete");
+        var results = DurableFuture.allOf(orderFuture, inventoryFuture, shippingFuture);
+
+        // Combine into summary
+        var summary = String.join(" | ", results);
+        context.getLogger().info("All child contexts complete: {}", summary);
+
+        return summary;
+    }
+}

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

@@ -530,6 +530,23 @@ void testChildContextExample() {
         assertNotNull(runner.getOperation("shipping-estimate"));
     }
 
+    @Test
+    void testVirtualChildContextExample() {
+        var runner = CloudDurableTestRunner.create(
+                arn("virtual-child-context-example"), GreetingRequest.class, String.class, lambdaClient);
+        var result = runner.run(new GreetingRequest("Alice"));
+
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+        assertEquals(
+                "Order for Alice [validated] | Stock available for Alice [confirmed] | Base rate for Alice + regional adjustment [shipping ready]",
+                result.getResult());
+
+        // Verify step operations in child context were tracked
+        assertNotNull(runner.getOperation("validate-order"));
+        assertNotNull(runner.getOperation("confirm-inventory"));
+        assertNotNull(runner.getOperation("finalize-shipping"));
+    }
+
     @ParameterizedTest
     @CsvSource({"100, 1000, 20", "500, 2000, 30", "1000, 3000, 50"})
     void testManyAsyncStepsExample(int steps, long maxExecutionTime, long maxReplayTime) {

examples/src/test/java/software/amazon/lambda/durable/examples/child/VirtualChildContextExampleTest.java

@@ -0,0 +1,56 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable.examples.child;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import org.junit.jupiter.api.Test;
+import software.amazon.lambda.durable.examples.types.GreetingRequest;
+import software.amazon.lambda.durable.model.ExecutionStatus;
+import software.amazon.lambda.durable.testing.LocalDurableTestRunner;
+
+class VirtualChildContextExampleTest {
+
+    @Test
+    void testVirtualChildContextExampleRunsToCompletion() {
+        var handler = new VirtualChildContextExample();
+        var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler);
+
+        var input = new GreetingRequest("Alice");
+        var result = runner.runUntilComplete(input);
+
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+        assertEquals(
+                "Order for Alice [validated] | Stock available for Alice [confirmed] | Base rate for Alice + regional adjustment [shipping ready]",
+                result.getResult(String.class));
+    }
+
+    @Test
+    void testVirtualChildContextExampleSuspendsOnFirstRun() {
+        var handler = new VirtualChildContextExample();
+        var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler);
+
+        var input = new GreetingRequest("Bob");
+
+        // First run should suspend due to wait operations inside child contexts
+        var result = runner.run(input);
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+    }
+
+    @Test
+    void testVirtualChildContextExampleReplay() {
+        var handler = new VirtualChildContextExample();
+        var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler);
+
+        var input = new GreetingRequest("Alice");
+
+        // First full execution
+        var result1 = runner.runUntilComplete(input);
+        assertEquals(ExecutionStatus.SUCCEEDED, result1.getStatus());
+
+        // Replay — should return cached results
+        var result2 = runner.run(input);
+        assertEquals(ExecutionStatus.SUCCEEDED, result2.getStatus());
+        assertEquals(result1.getResult(String.class), result2.getResult(String.class));
+    }
+}

examples/template.yaml

@@ -296,6 +296,23 @@ Resources:
                 - lambda:GetDurableExecutionState
               Resource: !Sub "arn:aws:lambda:${AWS::Region}:${AWS::AccountId}:function:child-context-example-${JavaVersion}-runtime"
 
+  VirtualChildContextExampleFunction:
+    Type: AWS::Serverless::Function
+    Properties:
+      FunctionName: !Join
+        - '-'
+        - - 'virtual-child-context-example'
+          - !Ref JavaVersion
+          - runtime
+      Handler: "software.amazon.lambda.durable.examples.child.VirtualChildContextExample"
+      Policies:
+        - Statement:
+            - Effect: Allow
+              Action:
+                - lambda:CheckpointDurableExecutions
+                - lambda:GetDurableExecutionState
+              Resource: !Sub "arn:aws:lambda:${AWS::Region}:${AWS::AccountId}:function:virtual-child-context-example-${JavaVersion}-runtime"
+
   WaitAsyncExampleFunction:
     Type: AWS::Serverless::Function
     Properties:

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

@@ -8,6 +8,7 @@
 import java.util.concurrent.atomic.AtomicInteger;
 import org.junit.jupiter.api.Test;
 import software.amazon.awssdk.services.lambda.model.OperationType;
+import software.amazon.lambda.durable.config.RunInChildContextConfig;
 import software.amazon.lambda.durable.model.ExecutionStatus;
 import software.amazon.lambda.durable.testing.LocalDurableTestRunner;
 
@@ -42,6 +43,34 @@ void childContextResultSurvivesReplay() {
         assertEquals(1, childExecutionCount.get(), "Child function should not re-execute on replay");
     }
 
+    @Test
+    void virtualChildContextResultSurvivesReplay() {
+        var childExecutionCount = new AtomicInteger(0);
+
+        var runner = LocalDurableTestRunner.create(
+                String.class,
+                (input, ctx) -> ctx.runInChildContext(
+                        "compute",
+                        TypeToken.get(String.class),
+                        child -> {
+                            childExecutionCount.incrementAndGet();
+                            return child.step("work", String.class, stepCtx -> "result-" + input);
+                        },
+                        RunInChildContextConfig.builder().isVirtual(true).build()));
+
+        // First run - executes child context
+        var result = runner.runUntilComplete("test");
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+        assertEquals("result-test", result.getResult(String.class));
+        assertEquals(1, childExecutionCount.get());
+
+        // Second run - replays, should return cached result without re-executing
+        result = runner.run("test");
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+        assertEquals("result-test", result.getResult(String.class));
+        assertEquals(2, childExecutionCount.get(), "Child function should re-execute on replay");
+    }
+
     /**
      * A child context that fails with a reconstructable exception SHALL preserve the exception type, message, and error
      * details through the checkpoint-and-replay cycle.
@@ -72,6 +101,36 @@ void childContextExceptionPreservedOnReplay() {
         assertEquals(1, childExecutionCount.get(), "Child function should not re-execute on failed replay");
     }
 
+    @Test
+    void virtualChildContextExceptionPreservedOnReplay() {
+        var childExecutionCount = new AtomicInteger(0);
+
+        var runner = LocalDurableTestRunner.create(
+                String.class,
+                (input, ctx) -> ctx.runInChildContext(
+                        "failing",
+                        String.class,
+                        child -> {
+                            childExecutionCount.incrementAndGet();
+                            throw new IllegalArgumentException("bad input: " + input);
+                        },
+                        RunInChildContextConfig.builder().isVirtual(true).build()));
+
+        // First run - child context fails
+        var result = runner.run("test");
+        assertEquals(ExecutionStatus.FAILED, result.getStatus());
+        assertEquals(1, childExecutionCount.get());
+
+        // Second run - replays, should throw same exception without re-executing
+        result = runner.run("test");
+        assertEquals(ExecutionStatus.FAILED, result.getStatus());
+        assertTrue(result.getError().isPresent());
+        var error = result.getError().get();
+        assertEquals("java.lang.IllegalArgumentException", error.errorType());
+        assertEquals("bad input: test", error.errorMessage());
+        assertEquals(2, childExecutionCount.get(), "Child function should re-execute on failed replay");
+    }
+
     /** Operations checkpointed from within a child context SHALL have the child context's ID as their parentId. */
     @Test
     void operationsInChildContextHaveCorrectParentId() {

sdk/src/test/java/software/amazon/lambda/durable/operation/ChildContextOperationTest.java

@@ -66,6 +66,15 @@ private ChildContextOperation<String> createOperation(Function<DurableContext, S
                 durableContext);
     }
 
+    private ChildContextOperation<String> createVirtualOperation(Function<DurableContext, String> func) {
+        return new ChildContextOperation<>(
+                OPERATION_IDENTIFIER,
+                func,
+                TypeToken.get(String.class),
+                RunInChildContextConfig.builder().serDes(SERDES).isVirtual(true).build(),
+                durableContext);
+    }
+
     private ChildContextOperation<String> createOperationWithParent(
             Function<DurableContext, String> func, ConcurrencyOperation<?> parent) {
         return new ChildContextOperation<>(
@@ -107,6 +116,22 @@ void replaySucceededReturnsCachedResult() {
         assertFalse(functionCalled.get(), "Function should not be called during SUCCEEDED replay");
     }
 
+    /** Virtual contexts are always executed, even during SUCCEEDED replay. */
+    @Test
+    void executeVirtualContext() {
+        var functionCalled = new AtomicBoolean(false);
+        var operation = createVirtualOperation(ctx -> {
+            functionCalled.set(true);
+            return "should-execute";
+        });
+
+        operation.execute();
+        var result = operation.get();
+
+        assertEquals("should-execute", result);
+        assertTrue(functionCalled.get(), "Function should be called during SUCCEEDED replay");
+    }
+
     // ===== FAILED replay =====
 
     /** FAILED replay throws the original exception without re-executing. */