Making assertions in tests against queries at various, non-sleep points in time

Take the following sample workflow.

class MyWorkflow {
  Activities activities;
  int i = 0;
  
  @WorkflowMethod
  void workflow() {
    activities.activityOne();
    i++; // i = 1;
    Workflow.sleep(Duration.ofMinutes(1));
    activities.activityTwo();
    i++; // i = 2;
    activities.activityThree();
    i++; // i = 3;
  }
  
  @QueryMethod
  int counter() {
    return i;
  }
}

I know after initially invoking the workflow, it will reach the sleep call, and I can assert that the result of counter query is 1. I can then call TestEnv.sleep, however then the workflow continues and the result of counter query is 3.

Is there a way to hold up the workflow at the call to activityThree() and make assertions that the counter query returns 2? I know this is trivial but gets at my desired to assert against workflow state at various points along its execution. I can do it with timers/sleep calls at the moment, but not other awaitables that the workflow would actually encounter.

(If it’s important, my activites are mocked for this test, I’m just testing workflow logic in isolation here.)

Phrasing my query a slightly different way, I found this answer from Maxim:

You can make queries from the mocked activities. This way you can verify that at the point that activity is invoked the query returns the expected value.

This definitely works and is probably what I will do for now barring any other answers, but it does feel funky. For one, it creates a really strange legibility. The assertion doesn’t “live” at the point of the test where it runs, but instead in the mocked activity closure. And while I understand that querying from an activity is logically the same as querying from the “outside world”, I don’t know how clear that would be to future maintainers or people on my team who read this test.

You could use workflow interceptor in test, something like this note you can run this interceptor in test only.
It’s similar solution as to querying from mocked activities, but does it from interceptor

The mocked activity can unblock the test.

For example; sdk-java/temporal-sdk/src/test/java/io/temporal/workflow/versionTests/GetVersionWithoutCommandEventTest.java at ce699764bf4a84dbf8abf4fdfca2eb4ff8344230 · temporalio/sdk-java · GitHub. The unit test unblocks the Future from the workflow, but in your case it can be done from the mocked activity.