Deflake the DefaultControllerTest.

Author: brendandburns

extended/src/test/java/io/kubernetes/client/extended/controller/DefaultControllerTest.java

@@ -25,13 +25,16 @@
 import java.util.List;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.concurrent.Semaphore;
 import java.util.concurrent.atomic.AtomicBoolean;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
+import org.mockito.invocation.InvocationOnMock;
 import org.mockito.junit.MockitoJUnitRunner;
+import org.mockito.stubbing.Answer;
 
 @RunWith(MockitoJUnitRunner.class)
 public class DefaultControllerTest {
@@ -40,16 +43,6 @@ public class DefaultControllerTest {
   private RateLimitingQueue<Request> workQueue =
       new DefaultRateLimitingQueue<>(Executors.newSingleThreadExecutor());
 
-  private final int stepCooldownIntervalInMillis = 500;
-
-  private void cooldown() {
-    try {
-      Thread.sleep(stepCooldownIntervalInMillis);
-    } catch (InterruptedException e) {
-      e.printStackTrace();
-    }
-  }
-
   @Before
   public void setUp() throws Exception {}
 
@@ -67,33 +60,53 @@ public void testStartingStoppingController() throws InterruptedException {
     testController.setWorkerThreadPool(Executors.newScheduledThreadPool(1));
 
     Request request1 = new Request("test1");
-    when(mockReconciler.reconcile(request1)).thenReturn(new Result(false));
+    final Semaphore latch = new Semaphore(1);
+    latch.acquire();
+    when(mockReconciler.reconcile(request1))
+        .thenAnswer(
+            new Answer() {
+              public Object answer(InvocationOnMock invocation) {
+                latch.release();
+                return new Result(false);
+              }
+            });
 
     // emit an event when the controller hasn't started
     workQueue.add(request1);
-    cooldown();
-    verify(mockReconciler, times(0)).reconcile(request1);
+    latch.acquire();
 
+    verify(mockReconciler, times(0)).reconcile(request1);
     controllerThead.submit(testController::run);
 
-    cooldown();
+    latch.acquire();
     verify(mockReconciler, times(1)).reconcile(request1);
 
     testController.shutdown();
     Request request2 = new Request("test2");
 
     // emit an event after the controller has shutdown
     workQueue.add(request2);
-    cooldown();
+    latch.acquire();
+
     verify(mockReconciler, times(0)).reconcile(request2);
   }
 
   @Test
   public void testControllerWontStartBeforeReady() throws InterruptedException {
 
     Request request1 = new Request("test1");
-    when(mockReconciler.reconcile(request1)).thenReturn(new Result(false));
+    final Semaphore latch = new Semaphore(1);
+
+    when(mockReconciler.reconcile(request1))
+        .thenAnswer(
+            new Answer() {
+              public Object answer(InvocationOnMock invocation) {
+                latch.release();
+                return new Result(false);
+              }
+            });
 
+    latch.acquire();
     AtomicBoolean ready = new AtomicBoolean(false);
     DefaultController testController =
         new DefaultController("", mockReconciler, workQueue, () -> ready.get());
@@ -105,12 +118,12 @@ public void testControllerWontStartBeforeReady() throws InterruptedException {
 
     // emit an event when the controller hasn't been ready
     workQueue.add(request1);
-    cooldown();
+    latch.acquire();
 
     verify(mockReconciler, times(0)).reconcile(request1);
 
     ready.set(true);
-    cooldown();
+    latch.acquire();
     verify(mockReconciler, times(1)).reconcile(request1);
   }
 
@@ -120,18 +133,23 @@ public void testControllerKeepsWorkingWhenReconcilerAbortsWithRuntimeException()
     AtomicBoolean aborts = new AtomicBoolean(true);
     AtomicBoolean resumed = new AtomicBoolean(false);
     List<Request> finishedRequests = new ArrayList<>();
+    final Semaphore latch = new Semaphore(1);
     DefaultController testController =
         new DefaultController(
             "",
             new Reconciler() {
               @Override
               public Result reconcile(Request request) {
-                if (aborts.get()) {
-                  throw new RuntimeException("Oops!!");
+                try {
+                  if (aborts.get()) {
+                    throw new RuntimeException("Oops!!");
+                  }
+                  resumed.set(true);
+                  finishedRequests.add(request);
+                  return new Result(false);
+                } finally {
+                  latch.release();
                 }
-                resumed.set(true);
-                finishedRequests.add(request);
-                return new Result(false);
               }
             },
             workQueue);
@@ -142,13 +160,13 @@ public Result reconcile(Request request) {
 
     Request request1 = new Request("test1");
     workQueue.add(request1);
-    cooldown();
+    latch.acquire();
 
     aborts.set(false);
     // emit another event, the previous one has been backoff'd
     Request request2 = new Request("test2");
     workQueue.add(request2);
-    cooldown();
+    latch.acquire();
     testController.shutdown();
 
     assertTrue(resumed.get());