Added Workflow.addWorkflowImplementationFactory for unit testing (#140)

* Added Worker.addWorkflowImplementationFactory

* Added unit test for worker.addWorkflowImplementationFactory

Author: mfateev

src/main/java/com/uber/cadence/internal/sync/POJOWorkflowImplementationFactory.java

@@ -28,6 +28,7 @@
 import com.uber.cadence.internal.worker.WorkflowExecutionException;
 import com.uber.cadence.testing.SimulatedTimeoutException;
 import com.uber.cadence.workflow.Functions;
+import com.uber.cadence.workflow.Functions.Func;
 import com.uber.cadence.workflow.QueryMethod;
 import com.uber.cadence.workflow.SignalMethod;
 import com.uber.cadence.workflow.Workflow;
@@ -59,9 +60,12 @@ final class POJOWorkflowImplementationFactory implements ReplayWorkflowFactory {
   private final Map<String, Functions.Func<SyncWorkflowDefinition>> workflowDefinitions =
       Collections.synchronizedMap(new HashMap<>());
 
+  private final Map<Class<?>, Functions.Func<?>> workflowImplementationFactories =
+      Collections.synchronizedMap(new HashMap<>());
+
   private final ExecutorService threadPool;
 
-  public POJOWorkflowImplementationFactory(
+  POJOWorkflowImplementationFactory(
       DataConverter dataConverter,
       ExecutorService threadPool,
       Function<WorkflowInterceptor, WorkflowInterceptor> interceptorFactory) {
@@ -70,14 +74,19 @@ public POJOWorkflowImplementationFactory(
     this.interceptorFactory = Objects.requireNonNull(interceptorFactory);
   }
 
-  public void setWorkflowImplementationTypes(Class<?>[] workflowImplementationTypes) {
+  void setWorkflowImplementationTypes(Class<?>[] workflowImplementationTypes) {
     workflowDefinitions.clear();
     for (Class<?> type : workflowImplementationTypes) {
       addWorkflowImplementationType(type);
     }
   }
 
-  public void addWorkflowImplementationType(Class<?> workflowImplementationClass) {
+  <R> void addWorkflowImplementationFactory(Class<R> clazz, Functions.Func<R> factory) {
+    workflowImplementationFactories.put(clazz, factory);
+    addWorkflowImplementationType(clazz);
+  }
+
+  private void addWorkflowImplementationType(Class<?> workflowImplementationClass) {
     TypeToken<?>.TypeSet interfaces =
         TypeToken.of(workflowImplementationClass).getTypes().interfaces();
     if (interfaces.isEmpty()) {
@@ -230,17 +239,22 @@ public byte[] execute(byte[] input) throws CancellationException, WorkflowExecut
 
     private void newInstance() {
       if (workflow == null) {
-        try {
-          workflow = workflowImplementationClass.getDeclaredConstructor().newInstance();
-        } catch (NoSuchMethodException
-            | InstantiationException
-            | IllegalAccessException
-            | InvocationTargetException e) {
-          // Error to fail decision as this can be fixed by a new deployment.
-          throw new Error(
-              "Failure instantiating workflow implementation class "
-                  + workflowImplementationClass.getName(),
-              e);
+        Func<?> factory = workflowImplementationFactories.get(workflowImplementationClass);
+        if (factory != null) {
+          workflow = factory.apply();
+        } else {
+          try {
+            workflow = workflowImplementationClass.getDeclaredConstructor().newInstance();
+          } catch (NoSuchMethodException
+              | InstantiationException
+              | IllegalAccessException
+              | InvocationTargetException e) {
+            // Error to fail decision as this can be fixed by a new deployment.
+            throw new Error(
+                "Failure instantiating workflow implementation class "
+                    + workflowImplementationClass.getName(),
+                e);
+          }
         }
         WorkflowInternal.registerQuery(workflow);
       }
@@ -289,7 +303,7 @@ public void processSignal(String signalName, byte[] input, long eventId) {
     }
   }
 
-  public static WorkflowExecutionException mapToWorkflowExecutionException(
+  static WorkflowExecutionException mapToWorkflowExecutionException(
       Exception failure, DataConverter dataConverter) {
     failure = CheckedExceptionWrapper.unwrap(failure);
     // Only expected during unit tests.

src/main/java/com/uber/cadence/internal/sync/SyncWorkflowWorker.java

@@ -23,6 +23,7 @@
 import com.uber.cadence.internal.worker.SingleWorkerOptions;
 import com.uber.cadence.internal.worker.WorkflowWorker;
 import com.uber.cadence.serviceclient.IWorkflowService;
+import com.uber.cadence.workflow.Functions.Func;
 import com.uber.cadence.workflow.WorkflowInterceptor;
 import java.util.concurrent.SynchronousQueue;
 import java.util.concurrent.ThreadPoolExecutor;
@@ -63,6 +64,10 @@ public void setWorkflowImplementationTypes(Class<?>[] workflowImplementationType
     factory.setWorkflowImplementationTypes(workflowImplementationTypes);
   }
 
+  public <R> void addWorkflowImplementationFactory(Class<R> clazz, Func<R> factory) {
+    this.factory.addWorkflowImplementationFactory(clazz, factory);
+  }
+
   public void start() {
     worker.start();
   }

src/main/java/com/uber/cadence/worker/Worker.java

@@ -17,13 +17,15 @@
 
 package com.uber.cadence.worker;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.uber.cadence.WorkflowExecution;
 import com.uber.cadence.client.WorkflowClient;
 import com.uber.cadence.internal.sync.SyncActivityWorker;
 import com.uber.cadence.internal.sync.SyncWorkflowWorker;
 import com.uber.cadence.internal.worker.SingleWorkerOptions;
 import com.uber.cadence.serviceclient.IWorkflowService;
 import com.uber.cadence.serviceclient.WorkflowServiceTChannel;
+import com.uber.cadence.workflow.Functions.Func;
 import java.time.Duration;
 import java.util.Objects;
 import java.util.concurrent.TimeUnit;
@@ -175,6 +177,32 @@ public void registerWorkflowImplementationTypes(Class<?>... workflowImplementati
     workflowWorker.setWorkflowImplementationTypes(workflowImplementationClasses);
   }
 
+  /**
+   * Configures a factory to use when an instance of a workflow implementation is created. The only
+   * valid use for this method is unit testing, specifically to instantiate mocks that implement
+   * child workflows. An example of mocking a child workflow:
+   *
+   * <pre><code>
+   *   worker.addWorkflowImplementationFactory(ChildWorkflow.class, () -> {
+   *     ChildWorkflow child = mock(ChildWorkflow.class);
+   *     when(child.workflow(anyString(), anyString())).thenReturn("result1");
+   *     return child;
+   *   });
+   * </code></pre>
+   *
+   * <p>Unless mocking a workflow execution use {@link
+   * #registerWorkflowImplementationTypes(Class[])}.
+   *
+   * @param workflowInterface Workflow interface that this factory implements
+   * @param factory factory that when called creates a new instance of the workflow implementation
+   *     object.
+   * @param <R> type of the workflow object to create.
+   */
+  @VisibleForTesting
+  public <R> void addWorkflowImplementationFactory(Class<R> workflowInterface, Func<R> factory) {
+    workflowWorker.addWorkflowImplementationFactory(workflowInterface, factory);
+  }
+
   /**
    * Register activity implementation objects with a worker. Overwrites previously registered
    * objects. As activities are reentrant and stateless only one instance per activity type is

src/test/java/com/uber/cadence/internal/testing/WorkflowTestingTest.java

@@ -20,6 +20,9 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
+import static org.mockito.Matchers.anyString;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 import com.uber.cadence.EventType;
 import com.uber.cadence.GetWorkflowExecutionHistoryRequest;
@@ -660,4 +663,27 @@ public void testChildSimulatedTimeout() {
       assertTrue(e.getCause() instanceof ChildWorkflowTimedOutException);
     }
   }
+
+  @Test
+  public void testMockedChildSimulatedTimeout() {
+    Worker worker = testEnvironment.newWorker(TASK_LIST);
+    worker.registerWorkflowImplementationTypes(SimulatedTimeoutParentWorkflow.class);
+    worker.addWorkflowImplementationFactory(
+        ChildWorkflow.class,
+        () -> {
+          ChildWorkflow child = mock(ChildWorkflow.class);
+          when(child.workflow(anyString(), anyString())).thenThrow(new SimulatedTimeoutException());
+          return child;
+        });
+    worker.start();
+    WorkflowClient client = testEnvironment.newWorkflowClient();
+    WorkflowOptions options = new WorkflowOptions.Builder().setWorkflowId("parent1").build();
+    ParentWorkflow workflow = client.newWorkflowStub(ParentWorkflow.class, options);
+    try {
+      workflow.workflow("input1");
+      fail("unreacheable");
+    } catch (WorkflowException e) {
+      assertTrue(e.getCause() instanceof ChildWorkflowTimedOutException);
+    }
+  }
 }