From 4dfd74906c50d69a94658fd2c03300e2e4abd94f Mon Sep 17 00:00:00 2001
From: Ravi Khadiwala <ravi@signal.org>
Date: Tue, 30 Jan 2024 13:10:43 -0600
Subject: [PATCH] Add timeouts to VirtualThreadPinEventMonitorTest

---
 .../util/VirtualThreadPinEventMonitorTest.java    | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/util/VirtualThreadPinEventMonitorTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/util/VirtualThreadPinEventMonitorTest.java
index 4c93704eb..209a9d5b5 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/util/VirtualThreadPinEventMonitorTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/util/VirtualThreadPinEventMonitorTest.java
@@ -53,8 +53,9 @@ public class VirtualThreadPinEventMonitorTest {
     exec.submit(() -> synchronizedSleep1()).get();
     eventMonitor.stop();
 
-    final Pair<RecordedEvent, Boolean> nxt = bq.take();
-    assertThat(nxt.getRight()).isFalse();
+    final Pair<RecordedEvent, Boolean> event = bq.poll(1, TimeUnit.SECONDS);
+    assertThat(event).isNotNull();
+    assertThat(event.getRight()).isFalse();
     assertThat(bq.isEmpty());
     exec.shutdown();
     exec.awaitTermination(1, TimeUnit.MILLISECONDS);
@@ -73,10 +74,12 @@ public class VirtualThreadPinEventMonitorTest {
     exec.submit(() -> synchronizedSleep2()).get();
     eventMonitor.stop();
 
-    final boolean sleep1Allowed = bq.take().getRight();
-    final boolean sleep2Allowed = bq.take().getRight();
-    assertThat(sleep1Allowed).isTrue();
-    assertThat(sleep2Allowed).isFalse();
+    final Pair<RecordedEvent, Boolean> sleep1Event = bq.poll(1, TimeUnit.SECONDS);
+    final Pair<RecordedEvent, Boolean> sleep2Event = bq.poll(1, TimeUnit.SECONDS);
+    assertThat(sleep1Event).isNotNull();
+    assertThat(sleep2Event).isNotNull();
+    assertThat(sleep1Event.getRight()).isTrue();
+    assertThat(sleep2Event.getRight()).isFalse();
     assertThat(bq.isEmpty());
     exec.shutdown();
     exec.awaitTermination(1, TimeUnit.MILLISECONDS);