Fix unit tests

Author: slinkydeveloper

sdk-core/src/main/java/dev/restate/sdk/core/HandlerContextImpl.java

@@ -416,10 +416,8 @@ private void pollAsyncResultInner(AsyncResultInternal<?> asyncResult) {
       } else if (response instanceof StateMachine.DoProgressResponse.ReadFromInput
           || response instanceof StateMachine.DoProgressResponse.WaitingPendingRun) {
         this.stateMachine.onNextEvent(
-            () -> {
-              LOG.info("Triggered after wait {}", response);
-              this.pollAsyncResultInner(asyncResult);
-            });
+            () -> this.pollAsyncResultInner(asyncResult),
+            response instanceof StateMachine.DoProgressResponse.ReadFromInput);
         return;
       } else if (response instanceof StateMachine.DoProgressResponse.ExecuteRun) {
         triggerScheduledRun(((StateMachine.DoProgressResponse.ExecuteRun) response).handle());

sdk-core/src/main/java/dev/restate/sdk/core/statemachine/StateMachine.java

@@ -44,7 +44,7 @@ static StateMachine init(
 
   // --- Await next event
 
-  void onNextEvent(Runnable runnable);
+  void onNextEvent(Runnable runnable, boolean triggerNowIfInputClosed);
 
   // --- Async results
 

sdk-core/src/main/java/dev/restate/sdk/core/statemachine/StateMachineImpl.java

@@ -73,14 +73,14 @@ public CompletableFuture<Void> waitForReady() {
   }
 
   @Override
-  public void onNextEvent(Runnable runnable) {
+  public void onNextEvent(Runnable runnable, boolean triggerNowIfInputClosed) {
     this.nextEventListener =
         () -> {
           this.nextEventListener.run();
           runnable.run();
         };
     // Trigger this now
-    if (this.stateContext.isInputClosed()) {
+    if (triggerNowIfInputClosed && this.stateContext.isInputClosed()) {
       this.triggerNextEventSignal();
     }
   }