Handle case where getMemo requests a non-existent key

temporal-sdk/src/main/java/io/temporal/internal/sync/WorkflowInternal.java

@@ -431,6 +431,10 @@ public static WorkflowInfo getWorkflowInfo() {
 
   public static <T> T getMemo(String key, Class<T> valueClass, Type valueType) {
     Payload memo = getRootWorkflowContext().getContext().getMemo(key);
+    if (memo == null) {
+      return null;
+    }
+
     return getDataConverter().fromPayload(memo, valueClass, valueType);
   }
 

temporal-sdk/src/main/java/io/temporal/workflow/Workflow.java

@@ -621,7 +621,7 @@ public static WorkflowInfo getInfo() {
    *
    * @param key memo key
    * @param valueClass Java class to deserialize into
-   * @return deserialized Memo
+   * @return deserialized Memo or null if the key is not present in the memo
    */
   public static <T> Object getMemo(String key, Class<T> valueClass) {
     return getMemo(key, valueClass, valueClass);
@@ -636,7 +636,7 @@ public static <T> Object getMemo(String key, Class<T> valueClass) {
    * @param key memo key
    * @param valueClass Java class to deserialize into
    * @param genericType type parameter for the generic class
-   * @return deserialized Memo
+   * @return deserialized Memo or null if the key is not present in the memo
    */
   public static <T> T getMemo(String key, Class<T> valueClass, Type genericType) {
     return WorkflowInternal.getMemo(key, valueClass, genericType);

temporal-sdk/src/test/java/io/temporal/workflow/MemoTest.java

@@ -19,8 +19,7 @@
 
 package io.temporal.workflow;
 
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.*;
 
 import com.google.common.collect.ImmutableMap;
 import com.google.common.reflect.TypeToken;
@@ -33,11 +32,11 @@
 import io.temporal.api.workflowservice.v1.GetWorkflowExecutionHistoryResponse;
 import io.temporal.client.WorkflowClient;
 import io.temporal.client.WorkflowOptions;
+import io.temporal.common.converter.DataConverterException;
 import io.temporal.common.converter.GsonJsonPayloadConverter;
 import io.temporal.internal.client.WorkflowClientHelper;
 import io.temporal.testing.internal.SDKTestOptions;
 import io.temporal.testing.internal.SDKTestWorkflowRule;
-import io.temporal.workflow.shared.TestMultiArgWorkflowFunctions.TestMultiArgWorkflowImpl;
 import io.temporal.workflow.shared.TestMultiArgWorkflowFunctions.TestNoArgsWorkflowFunc;
 import io.temporal.workflow.shared.TestWorkflows.NoArgsWorkflow;
 import java.util.HashMap;
@@ -51,6 +50,7 @@ public class MemoTest {
   private static final String MEMO_VALUE = "testValue";
   private static final String MEMO_KEY_2 = "testKey2";
   private static final Integer MEMO_VALUE_2 = 1;
+  private static final String MEMO_KEY_MISSING = "testKey3";
   private static final Map<String, Object> MEMO =
       ImmutableMap.of(
           MEMO_KEY,
@@ -64,9 +64,11 @@ public class MemoTest {
 
   @Rule
   public SDKTestWorkflowRule testWorkflowRule =
-      SDKTestWorkflowRule.newBuilder()
-          .setWorkflowTypes(TestMultiArgWorkflowImpl.class, WorkflowImpl.class)
-          .build();
+      SDKTestWorkflowRule.newBuilder().setWorkflowTypes(WorkflowWithMemoImpl.class).build();
+
+  @Rule
+  public SDKTestWorkflowRule testNoMemoWorkflowRule =
+      SDKTestWorkflowRule.newBuilder().setWorkflowTypes(WorkflowWithoutMemoImpl.class).build();
 
   @Test
   public void testMemo() {
@@ -116,15 +118,54 @@ public void testMemoInWorkflow() {
     workflow.execute();
   }
 
-  public static class WorkflowImpl implements NoArgsWorkflow {
+  @Test
+  public void testNoMemoInWorkflowFailsGetMemoGracefully() {
+    WorkflowOptions workflowOptions =
+        SDKTestOptions.newWorkflowOptionsWithTimeouts(testNoMemoWorkflowRule.getTaskQueue())
+            .toBuilder()
+            .setTaskQueue(testNoMemoWorkflowRule.getTaskQueue())
+            .build();
+
+    NoArgsWorkflow workflow =
+        testNoMemoWorkflowRule
+            .getWorkflowClient()
+            .newWorkflowStub(NoArgsWorkflow.class, workflowOptions);
+    workflow.execute();
+  }
+
+  public static class WorkflowWithMemoImpl implements NoArgsWorkflow {
     @Override
     public void execute() {
+      // Simple value found
       assertEquals(MEMO_VALUE, Workflow.getMemo(MEMO_KEY, String.class));
+
+      // Map value found
       Map result =
           Workflow.getMemo(
               MEMO_KEY_2, Map.class, new TypeToken<HashMap<String, Integer>>() {}.getType());
       assertTrue(result instanceof HashMap);
       assertEquals(MEMO_VALUE_2, result.get(MEMO_KEY_2));
+
+      // Requested mismatched type
+      boolean throwsDataConverterException = false;
+      try {
+        Workflow.getMemo(MEMO_KEY, Integer.class);
+      } catch (DataConverterException e) {
+        throwsDataConverterException = true;
+      } finally {
+        assertTrue(throwsDataConverterException);
+      }
+
+      // Missing key not found
+      assertNull(Workflow.getMemo(MEMO_KEY_MISSING, String.class));
+    }
+  }
+
+  public static class WorkflowWithoutMemoImpl implements NoArgsWorkflow {
+    @Override
+    public void execute() {
+      // Missing key not found
+      assertNull(Workflow.getMemo(MEMO_KEY, String.class));
     }
   }
 }