remove Thread.sleep from test (#27)

tmio · Aug 22, 2023 · 465645c · 465645c
1 parent ae67213
commit 465645c
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/plumtree/src/test/java/org/apache/tuweni/plumtree/StateTest.java b/plumtree/src/test/java/org/apache/tuweni/plumtree/StateTest.java
@@ -8,6 +8,8 @@
 
 import org.apache.tuweni.bytes.Bytes;
 import org.apache.tuweni.bytes.Bytes32;
+import org.apache.tuweni.concurrent.AsyncCompletion;
+import org.apache.tuweni.concurrent.CompletableAsyncCompletion;
 import org.apache.tuweni.crypto.Hash;
 import org.apache.tuweni.junit.BouncyCastleExtension;
 
@@ -39,13 +41,18 @@ private static class MockMessageSender implements MessageSender {
     Bytes hash;
     Bytes payload;
 
+    CompletableAsyncCompletion waitForMessageSent;
+
     @Override
     public void sendMessage(
         Verb verb, Map<String, Bytes> attributes, Peer peer, Bytes hash, Bytes payload) {
       this.verb = verb;
       this.peer = peer;
       this.hash = hash;
       this.payload = payload;
+      if (waitForMessageSent != null) {
+        waitForMessageSent.complete();
+      }
     }
   }
 
@@ -257,6 +264,7 @@ void receiveFullMessageFromEagerPeerThenPartialMessageFromLazyPeer() {
   void receivePartialMessageFromLazyPeerAndNoFullMessage() throws Exception {
     EphemeralPeerRepository repo = new EphemeralPeerRepository();
     MockMessageSender messageSender = new MockMessageSender();
+    messageSender.waitForMessageSent = AsyncCompletion.incomplete();
     State state =
         new State(
             repo,
@@ -274,7 +282,7 @@ void receivePartialMessageFromLazyPeerAndNoFullMessage() throws Exception {
     repo.moveToLazy(lazyPeer);
     Bytes message = Bytes32.random();
     state.receiveIHaveMessage(lazyPeer, message);
-    Thread.sleep(200);
+    messageSender.waitForMessageSent.join();
     assertEquals(message, messageSender.hash);
     assertEquals(lazyPeer, messageSender.peer);
     assertEquals(MessageSender.Verb.GRAFT, messageSender.verb);