KAFKA-17871: fixup tests

Author: davide-armand

connect/runtime/src/test/java/org/apache/kafka/connect/storage/OffsetStorageReaderTest.java

@@ -28,8 +28,10 @@
 
 import java.util.Collections;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
 
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doAnswer;
@@ -57,9 +59,13 @@ public void testClosingOffsetReaderWhenOffsetStoreHangs() throws Exception {
         OffsetStorageReaderImpl offsetStorageReaderImpl = new OffsetStorageReaderImpl(
                 offsetBackingStore, "namespace", taskKeyConverter, taskValueConverter);
 
+        CountDownLatch latch = new CountDownLatch(1);
+
+        // Hanging `offsetBackingStore.get()`
         doAnswer(invocation -> {
-            // Sleep for a long time to simulate a hanging offset store
-            Thread.sleep(9999 * 1000);
+            latch.countDown();
+            CompletableFuture<Void> future = new CompletableFuture<>();
+            future.get(9999, TimeUnit.SECONDS);
             throw new RuntimeException("Should never get here");
         }).when(offsetBackingStore).get(any());
 
@@ -68,13 +74,20 @@ public void testClosingOffsetReaderWhenOffsetStoreHangs() throws Exception {
             // Does call offsetBackingStore.get() and hangs
             offsetStorageReaderImpl.offsets(Collections.emptyList());
         });
-        Thread.sleep(3000);
+
+        // Ensure the task is hanging
+        latch.await();
 
         verify(offsetBackingStore, times(1)).get(any());
 
         // The herder thread should not block when trying to close `offsetStorageReaderImpl`
         // and complete before test timeout
         offsetStorageReaderImpl.close();
+
+        executor.shutdownNow();
+        if (!executor.awaitTermination(5, TimeUnit.SECONDS)) {
+            throw new RuntimeException("Failed to shutdown executor");
+        }
     }
 
     @Test
@@ -89,9 +102,13 @@ public void testClosingOffsetReaderWhenOffsetStoreHangsAndHasIncompleteFutures()
         OffsetStorageReaderImpl offsetStorageReaderImpl = new OffsetStorageReaderImpl(
                 offsetBackingStore, "namespace", taskKeyConverter, taskValueConverter);
 
+        CountDownLatch latchTask1 = new CountDownLatch(1);
+        CountDownLatch latchTask2 = new CountDownLatch(1);
+
         // Mock hanging future
         doAnswer(invocation -> {
-            Thread.sleep(9999 * 1000);
+            CompletableFuture<Void> future = new CompletableFuture<>();
+            future.get(9999, TimeUnit.SECONDS);
             throw new RuntimeException("Should never get here");
         }).when(hangingFuture).get();
 
@@ -104,10 +121,13 @@ public Object answer(InvocationOnMock invocation) throws Throwable {
                     if (callCount == 0) {
                         callCount += 1;
                         // First connector task
+                        latchTask1.countDown();
                         return hangingFuture;
                     } else {
                         // Second connector task
-                        Thread.sleep(9999 * 1000);
+                        latchTask2.countDown();
+                        CompletableFuture<Void> future = new CompletableFuture<>();
+                        future.get(9999, TimeUnit.SECONDS);
                         throw new RuntimeException("Should never get here");
                     }
                 }
@@ -120,7 +140,8 @@ public Object answer(InvocationOnMock invocation) throws Throwable {
         executor.submit(() -> {
             offsetStorageReaderImpl.offsets(Collections.emptyList());
         });
-        Thread.sleep(3000);
+        // Ensure first task is hanging
+        latchTask1.await();
 
         verify(offsetBackingStore, times(1)).get(any());
         verify(hangingFuture, times(1)).get();
@@ -130,7 +151,8 @@ public Object answer(InvocationOnMock invocation) throws Throwable {
         executor.submit(() -> {
             offsetStorageReaderImpl.offsets(Collections.emptyList());
         });
-        Thread.sleep(3000);
+        // Ensure second task is hanging
+        latchTask2.await();
 
         verify(offsetBackingStore, times(2)).get(any());
 
@@ -140,5 +162,10 @@ public Object answer(InvocationOnMock invocation) throws Throwable {
 
         // The hanging future should be cancelled by `close()`
         verify(hangingFuture, times(1)).cancel(true);
+
+        executor.shutdownNow();
+        if (!executor.awaitTermination(5, TimeUnit.SECONDS)) {
+            throw new RuntimeException("Failed to shutdown executor");
+        }
     }
 }