diff --git a/temporal-sdk/src/test/java/io/temporal/internal/worker/AsyncPollerTest.java b/temporal-sdk/src/test/java/io/temporal/internal/worker/AsyncPollerTest.java index 5faa34ca7..22a6c2b49 100644 --- a/temporal-sdk/src/test/java/io/temporal/internal/worker/AsyncPollerTest.java +++ b/temporal-sdk/src/test/java/io/temporal/internal/worker/AsyncPollerTest.java @@ -336,7 +336,15 @@ public void testAsyncPollFailed() @Test public void testSuspendPolling() throws InterruptedException, ExecutionException, AsyncPoller.PollTaskAsyncAbort { - CountingSlotSupplier slotSupplierInner = new CountingSlotSupplier<>(1); + CountDownLatch reserveAttemptLatch = new CountDownLatch(2); + CountingSlotSupplier slotSupplierInner = + new CountingSlotSupplier(1) { + @Override + public SlotSupplierFuture reserveSlot(SlotReserveContext ctx) throws Exception { + reserveAttemptLatch.countDown(); + return super.reserveSlot(ctx); + } + }; TrackingSlotSupplier slotSupplier = new TrackingSlotSupplier<>(slotSupplierInner, new NoopScope()); DummyTaskExecutor executor = new DummyTaskExecutor(slotSupplier); @@ -377,6 +385,9 @@ public void testSuspendPolling() assertEquals(1, slotSupplierInner.reservedCount.get()); assertEquals(0, slotSupplier.getUsedSlots().size()); }); + // Wait until the poll loop is reserving a slot for its next iteration. Suspending before this + // point would correctly prevent that reservation, making the expected count racy. + assertTrue(reserveAttemptLatch.await(5, TimeUnit.SECONDS)); // Suspend polling again, this will not affect the already issued poll request poller.suspendPolling(); completePoll.get().apply();