Case Study

Let's first start with a sample workflow that we can use throughout the article.

We will start with a dummy workflow that calls a few Temporal activities while waiting for the user's interaction through signals. 

public class DummyWorkflowImpl implements DummyWorkflow {

    private final DummyActivity activity = Workflow.newActivityStub(DummyActivity.class, ActivityOptions.newBuilder()
                    .setTaskQueue("mock-static").setScheduleToCloseTimeout(Duration.ofMinutes(1)).build());

    private final Logger log = Workflow.getLogger(DummyWorkflowImpl.class);

    private boolean waiting1;

    @Override
    public void execute() {
        log.info("workflow started");

        activity.doSomething1();

        waiting1 = true;
        Workflow.await(() -> !waiting1);

        activity.doSomething2();
    }

    @Override
    public void dummySignal() {
        waiting1 = false;
    }
}

About Context Propagation

After a close monitoring of the context propagation happening through the lifetime of an execution of the above workflow definition, we were able to trace the flow in the below sequence diagram:
 

Eventhough we're using a single Temporal worker, we can see that the workflow execution was distributed among three threads (58, 59, and 61). This hints to the workers reusing the thread for doing several tasks. Since this tracing was done through the usage of a context propagator, we can deduce that the calls for propagating the context happens before every Temporal task (even if the worker or the thread has already been used before; this detail is important to understand)

Mockito's Mock Static

There are two problems with using Mockito's mockStatic here:

1. the stubbed behavior is local to the thread of execution; i.e. it will be lost upon switching threads
2. calling mockStatic twice on the same thread with the same target class fails, unless the reference returned from the first call is closed

To address these issues, we can use a context propagator as follows:

public class MockStaticContextPropagator implements ContextPropagator {


    // thread local field; to avoid threads affecting each other
    private final ThreadLocal<MockedStatic<Workflow>> workflowMockedStatic = new ThreadLocal<>(); 

    @Override
    public String getName() {
        return "MockStaticContextPropagator";
    }

    @Override
    public Map<String, Payload> serializeContext(Object context) {
        return Map.of();
    }

    @Override
    public Object deserializeContext(Map<String, Payload> header) {
        return new Object();
    }

    @Override
    public Object getCurrentContext() {
        return null;
    }

    @Override
    public void setCurrentContext(Object context) {
        // 1. checking if an existing reference for the mock static is not closed
        if (workflowMockedStatic.get() != null) {
            workflowMockedStatic.get().close();
            workflowMockedStatic.remove();
        }


        // 2. create a reference for the mock static and store it in a thread local field
        workflowMockedStatic.set(mockStatic(Workflow.class));
        // 3. do the needed stubbing
        workflowMockedStatic.get().when(() -> Workflow.getLogger(any(Class.class))).thenReturn(mock(Logger.class));
    }
}

Since the setCurrentContext is called before every Temporal task, we can use it, in our case, as an interceptor.

After solving the main issues, two minor problems still challenge our solution:

1. calling mockStatic on a class mocks the behavior of all the static methods in this class.
2. the above solution is not flexible enough to be easily used in tests

To address the first issue, we can create wrapper classes containing only the static methods we wish to mock, and use this class instead of the main class in our workflow.

Regarding the second point, we can introduce some flexibility to our solution as follows:

public class MockStaticContextPropagator<T> implements ContextPropagator {

    private final Class<T> targetClass;

    public MockStaticContextPropagator(Class<T> targetClass) {
        this.targetClass = targetClass;
    }

    private final ThreadLocal<MockedStatic<T>> workflowMockedStatic = new ThreadLocal<>();
    private final List<Stubbing> stubbings = new ArrayList<>();

    @Override
    public String getName() {
        return "MockStaticContextPropagator";
    }

    @Override
    public Map<String, Payload> serializeContext(Object context) {
        return Map.of();
    }

    @Override
    public Object deserializeContext(Map<String, Payload> header) {
        return new Object();
    }

    @Override
    public Object getCurrentContext() {
        return null;
    }

    @Override
    public void setCurrentContext(Object context) {
        if (workflowMockedStatic.get() != null) {
            workflowMockedStatic.get().close();
            workflowMockedStatic.remove();
        }

        workflowMockedStatic.set(mockStatic(targetClass));
        stubbings.forEach(s -> workflowMockedStatic.get().when(s.verification()).thenAnswer(s.answer()));
    }

    public void addStubbing(MockedStatic.Verification verification, Answer<?> answer) {
        stubbings.add(new Stubbing(verification, answer));
    }

    private record Stubbing(MockedStatic.Verification verification, Answer<?> answer) { }
}

The full code can be found here: https://github.com/mohammed-ezzedine/articles-temporal-mock-static