Add capability for waiting until verification complete

src/main/java/com/trickl/flux/websocket/ClosedWebSocketStepsBuilder.java

@@ -3,6 +3,7 @@
 import com.trickl.exceptions.StepVerifierException;
 import java.time.Duration;
 import java.util.Queue;
+import java.util.concurrent.CountDownLatch;
 import java.util.function.Supplier;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.java.Log;
@@ -24,12 +25,18 @@ public final class ClosedWebSocketStepsBuilder {
 
   private final Queue<Runnable> steps;
 
-  /** Wait for the server to be started and available. */
+  /** 
+   * Wait for the server to be started and available.
+   * @return The verifier builder
+   * */
   public ClosedWebSocketStepsBuilder thenWaitServerShutdown() {
     return thenWaitServerShutdown(Duration.ofSeconds(10));
   }
 
-  /** Wait for the server to be started and available. */
+  /** Wait for the server to be started and available.
+   * @param timeout How long to wait
+   * @return The verifier builder    
+   */
   public ClosedWebSocketStepsBuilder thenWaitServerShutdown(Duration timeout) {
     steps.add(() -> testWasShutdown(timeout));
 
@@ -45,12 +52,18 @@ protected void testWasShutdown(Duration timeout) {
   }
 
 
-  /** Wait for the server to be started and available. */
+  /**
+   *  Wait for the server to be started and available.
+   * @return The verifier builder
+   *  */
   public ClosedWebSocketStepsBuilder thenWaitServerStartThenUpgrade() {
     return thenWaitServerStartThenUpgrade(Duration.ofSeconds(10));
   }
 
-  /** Wait for the server to be started and available. */
+  /** Wait for the server to be started and available.
+   * @param timeout How long to wait
+   * @return The verifier builder    
+   */
   public ClosedWebSocketStepsBuilder thenWaitServerStartThenUpgrade(Duration timeout) {
     steps.add(
         () -> {
@@ -70,7 +83,10 @@ protected void testWasStarted(Duration timeout) {
     }    
   }
 
-  /** Expect the socket to be opened. */
+  /** Expect the socket to be opened. 
+   * @return The verifier builder
+   * 
+  */
   public OpenWebSocketStepsBuilder thenExpectOpen() {
     return thenExpectOpen(Duration.ofSeconds(10));
   }
@@ -79,6 +95,7 @@ public OpenWebSocketStepsBuilder thenExpectOpen() {
    * Expect the socket to be opened.
    *
    * @param timeout How long to wait
+   * @return The verifier builder
    */
   public OpenWebSocketStepsBuilder thenExpectOpen(Duration timeout) {
     steps.add(() -> testWasOpen(timeout));
@@ -95,13 +112,28 @@ protected void testWasOpen(Duration timeout) {
     }        
   }
 
-  /** Verify the steps ran. */
-  public void thenVerify() {
+  /** Perform an action.
+   * @param step the action to complete
+   * @return The verifier builder   
+   */
+  public ClosedWebSocketStepsBuilder then(Runnable step) {
+    steps.add(step);
+    return this;
+  }
+
+  /** Verify the steps ran. 
+   * @return The verification completion object
+   * 
+  */
+  public VerifierComplete thenVerify() {
+    CountDownLatch completeSignal = new CountDownLatch(1);
     scheduler.schedule(
         () -> {
           while (!steps.isEmpty()) {
             steps.remove().run();
           }
+          completeSignal.countDown();
         });
-  }
+    return new VerifierComplete(completeSignal);
+  }  
 }

src/main/java/com/trickl/flux/websocket/MockServerWithWebSocket.java

@@ -26,7 +26,9 @@ public class MockServerWithWebSocket {
   private final Duration waitStartTimeout = Duration.ofSeconds(10);
   Semaphore canStartSemaphore = new Semaphore(1);
 
-  /** Expect an open event. */
+  /** Expect an open event. 
+   * @return The verifier builder 
+  */
   public ClosedWebSocketStepsBuilder beginVerifier() {
     Scheduler scheduler = Schedulers.newParallel("verifier");
     Queue<Runnable> steps = new ConcurrentLinkedQueue<>();
@@ -58,7 +60,9 @@ public void start() {
     }
   }
 
-  /** Shutdown the mock web server. */
+  /** Shutdown the mock web server. 
+   * @throws IOException if the server cannot be shutdown
+  */
   public void shutdown() throws IOException {
     mockServer.shutdown();
     serverListener.onShutdown();

src/main/java/com/trickl/flux/websocket/OpenWebSocketStepsBuilder.java

@@ -35,6 +35,7 @@ public final class OpenWebSocketStepsBuilder {
    * Schedule sending a websocket message.
    *
    * @param payload The message payloadfto send
+   * @return The verifier builder
    */
   public OpenWebSocketStepsBuilder thenSend(final String payload) {
     steps.add(
@@ -51,6 +52,7 @@ public OpenWebSocketStepsBuilder thenSend(final String payload) {
    * Expect a message within a timeframe.
    *
    * @param body expected message body
+   * @return The verifier builder
    */
   public OpenWebSocketStepsBuilder thenExpectMessage(String body) {
     return thenExpectMessage(body, Duration.ofSeconds(10));
@@ -61,6 +63,7 @@ public OpenWebSocketStepsBuilder thenExpectMessage(String body) {
    *
    * @param body expected message body
    * @param timeout How long to wait
+   * @return The verifier builder
    */
   public OpenWebSocketStepsBuilder thenExpectMessage(String body, Duration timeout) {
     return thenExpectMessage(text -> text.equals(body), timeout);
@@ -70,6 +73,7 @@ public OpenWebSocketStepsBuilder thenExpectMessage(String body, Duration timeout
    * Expect a message within a timeframe.
    *
    * @param bodyPattern test for message content
+   * @return The verifier builder
    */
   public OpenWebSocketStepsBuilder thenExpectMessage(Pattern bodyPattern) {
     return thenExpectMessage(bodyPattern, Duration.ofSeconds(10));
@@ -80,6 +84,7 @@ public OpenWebSocketStepsBuilder thenExpectMessage(Pattern bodyPattern) {
    *
    * @param bodyPattern test for message content
    * @param timeout How long to wait
+   * @return The verifier builder
    */
   public OpenWebSocketStepsBuilder thenExpectMessage(Pattern bodyPattern, Duration timeout) {
     return thenExpectMessage(text -> bodyPattern.matcher(text).matches(), timeout);
@@ -89,6 +94,7 @@ public OpenWebSocketStepsBuilder thenExpectMessage(Pattern bodyPattern, Duration
    * Expect a message within a timeframe.
    *
    * @param bodyMatcher test for message content
+   * @return The verifier builder
    */
   public OpenWebSocketStepsBuilder thenExpectMessage(Predicate<String> bodyMatcher) {
     return thenExpectMessage(bodyMatcher, Duration.ofSeconds(10));
@@ -99,6 +105,7 @@ public OpenWebSocketStepsBuilder thenExpectMessage(Predicate<String> bodyMatcher
    *
    * @param bodyMatcher test for message content
    * @param timeout How long to wait
+   * @return The verifier builder
    */
   public OpenWebSocketStepsBuilder thenExpectMessage(
       Predicate<String> bodyMatcher, Duration timeout) {
@@ -122,6 +129,7 @@ protected void testWasMessage(Predicate<String> bodyMatcher, Duration timeout) {
    * Wait a period of time doing nothing.
    *
    * @param period how long to wait
+   * @return The verifier builder
    */
   public OpenWebSocketStepsBuilder thenWait(Duration period) {
     steps.add(
@@ -137,7 +145,10 @@ public OpenWebSocketStepsBuilder thenWait(Duration period) {
     return this;
   }
 
-  /** Expect the socket to be closing. */
+  /** Expect the socket to be closing.
+   * @return The verifier builder
+   * 
+   */
   public ClosedWebSocketStepsBuilder thenExpectClose() {
     return thenExpectClose(Duration.ofSeconds(10));
   }
@@ -146,6 +157,7 @@ public ClosedWebSocketStepsBuilder thenExpectClose() {
    * Expect the socket to be closed.
    *
    * @param timeout How long to wait
+   * @return The verifier builder
    */
   public ClosedWebSocketStepsBuilder thenExpectClose(Duration timeout) {
     steps.add(
@@ -173,7 +185,10 @@ public ClosedWebSocketStepsBuilder thenExpectClose(Duration timeout) {
         serverSupplier, mockWebServerListener, mockWebSocketListener, scheduler, steps);
   }
 
-  /** Close the connection. */
+  /** Close the connection. 
+   * @return The verifier builder
+   * 
+  */
   public ClosedWebSocketStepsBuilder thenClose() {
     steps.add(
         () -> {
@@ -187,4 +202,13 @@ public ClosedWebSocketStepsBuilder thenClose() {
     return new ClosedWebSocketStepsBuilder(
         serverSupplier, mockWebServerListener, mockWebSocketListener, scheduler, steps);
   }
+
+  /** Perform an action. 
+   * @param step the action to complete
+   * @return The verifier builder
+  */
+  public OpenWebSocketStepsBuilder then(Runnable step) {
+    steps.add(step);
+    return this;
+  }
 }

src/main/java/com/trickl/flux/websocket/VerifierComplete.java

@@ -0,0 +1,37 @@
+package com.trickl.flux.websocket;
+
+import java.time.Duration;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import lombok.RequiredArgsConstructor;
+import lombok.extern.java.Log;
+
+@Log
+@RequiredArgsConstructor
+public class VerifierComplete {
+
+  private final CountDownLatch completeSignal;
+
+  /**
+   * Block until complete.
+   * @param duration how long to wait
+   */
+  public void waitComplete(Duration duration) {
+    try {
+      completeSignal.await(duration.toMillis(), TimeUnit.MILLISECONDS);
+    } catch (InterruptedException ex) {
+      log.warning("Verifier interrupted");
+    }
+  }
+
+  /**
+   * Block until complete.
+   */
+  public void waitComplete() {
+    try {
+      completeSignal.await();
+    } catch (InterruptedException ex) {
+      log.warning("Verifier interrupted");
+    }
+  }
+}

src/test/java/com/trickl/flux/websocket/EchoWebSocketTest.java

@@ -1,10 +1,8 @@
 package com.trickl.flux.websocket;
 
 import com.trickl.flux.config.WebSocketConfiguration;
-
 import java.io.IOException;
 import java.time.Duration;
-
 import org.junit.jupiter.api.Test;
 import org.springframework.boot.test.context.SpringBootTest;
 import org.springframework.http.HttpHeaders;
@@ -21,7 +19,7 @@ public void testEchoWebSocket() throws IOException {
 
     MockServerWithWebSocket mockServer = new MockServerWithWebSocket();
 
-    mockServer.beginVerifier()
+    final VerifierComplete verifierComplete = mockServer.beginVerifier()
         .thenWaitServerStartThenUpgrade()
         .thenExpectOpen()
         .thenSend("MESSAGE 1")
@@ -46,5 +44,7 @@ public void testEchoWebSocket() throws IOException {
         .log("SESSION")
         .block(Duration.ofSeconds(60));
     mockServer.shutdown();
+
+    verifierComplete.waitComplete();
   }
 }