Mocking workflow within an activity in JUnit 5 tests not working as expected

Hello,

I’m using JUnit 5 to test an activity which spawns a temporal workflow. I’ve registered a TestWorkflowExtension with my activity, and attempted to test it. There are two issues I’m running into:

  1. The activity is registered using setActivityImplementation({activity object}), but the activity object has WorkflowClient injected into it for creating the workflow. Since the WorkflowClient is passed into the test method as a parameter for usage in JUnit 5, this runs into a chicken and egg problem with creating the activity object with the correct WorkflowClient object, but also needing to set the activity implementation to the object when registering the extension (before getting access to the WorkflowClient object). The temporary fix I have for this is simply a test method setWorkflowClient within the activity object that allows you to change the object to the correct one when the test starts.
  2. I try to register a mock version of my workflow using registerWorkflowImplementationFactory using the worker object provided in the test input parameters, but for some reason it does not use the mock. Notably, the same approach I took works for child workflows. My guess as to why this is happening is that the worker used to run the activity is not the same as the one to run the workflow, but I am not entirely sure that’s the case.

Attached is the relevant code:

// activity code:
    /** {@inheritDoc} */
    @Override
    public void startMyWorkflow(@NonNull final StartMyWorkflowActivityInput input) {
        final var workflowId = UUID.randomUUID() + "-" + input.record().recordId();

// create search attributes

        final var workflowOptions = WorkflowOptions.newBuilder()
                .setWorkflowId(workflowId)
                .setTypedSearchAttributes(searchAttributes)
                .setTaskQueue(serviceProperties.workflowConfig().taskQueue())
                .build();
        final var workflow = workflowClient.newWorkflowStub(MyWorkflow.class, workflowOptions);

// ensure workflow isn't running already

        WorkflowClient.start(workflow::run, input.flow(), workflowInput);
        }



// test code:
    /** Workflow service stubs */
    @NonNull
    private final WorkflowServiceStubs serviceStubs = mock(WorkflowServiceStubs.class);

    /** my workflow */
    @NonNull
    private final MyWorkflow myWorkflow = mock(MyWorkflow.class);

    /** blocking stub */
    @NonNull
    private final WorkflowServiceGrpc.WorkflowServiceBlockingStub blockingStub =
            mock(WorkflowServiceGrpc.WorkflowServiceBlockingStub.class);

    /** list workflow execution response */
    @NonNull
    private final ListWorkflowExecutionsResponse executionsResponse = mock(ListWorkflowExecutionsResponse.class);

    /** The start my workflow activity */
    @NonNull
    private final StartMyWorkflowActivity startMyWorkflowActivity =
            new StartMyWorkflowActivityImpl(serviceProperties, mock(WorkflowClient.class), serviceStubs);

    /** workflow test extension */
    @RegisterExtension
    @NonNull
    private final TestWorkflowExtension testWorkflowExtension = TestWorkflowExtension.newBuilder()
            .setActivityImplementations(startMyWorkflowActivity)
            .build();

    @Test
    void startMyWorkflow(
            @NonNull final TestWorkflowEnvironment testEnv,
            @NonNull final Worker worker,
            @NonNull final WorkflowClient workflowClient) {

        startMyWorkflowActivity.setWorkflowClient(workflowClient);

        when(myWorkflow.run(any(), any())).thenReturn("test");

        worker.registerWorkflowImplementationFactory(MyWorkflow.class, () -> myWorkflow);

        when(serviceStubs.blockingStub()).thenReturn(blockingStub);
        when(blockingStub.listWorkflowExecutions(any())).thenReturn(executionsResponse);
        when(executionsResponse.getExecutionsList()).thenReturn(List.of());

        testEnv.start();

        // validation code below
}


Any ideas on how to fix this? Let me know if any more info is needed.