Replace sleep with CountDownLatch in tests (#126)

* Replace sleep with CountDownLatch in tests

Relying on Thread.sleep(5) in ContextWrappingTests leads to assertion failures
that imply broken functionality, while simply the operation would not finish in
the given time. This change replaces Thread.sleep with CountDownLatch to wait
for the task to finish and throws a TimeoutException in case it does not.

Fixes #125.

* Added a missed usage

Author: Dariusz Jędrzejczyk

context-propagation/src/test/java/io/micrometer/context/ContextWrappingTests.java

@@ -17,6 +17,7 @@
 
 import java.util.Collections;
 import java.util.concurrent.Callable;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executor;
 import java.util.concurrent.ExecutorService;
@@ -51,7 +52,7 @@ void clear() {
     }
 
     @Test
-    void should_instrument_runnable() throws InterruptedException {
+    void should_instrument_runnable() throws InterruptedException, TimeoutException {
         StringThreadLocalHolder.setValue("hello");
         AtomicReference<String> valueInNewThread = new AtomicReference<>();
         Runnable runnable = runnable(valueInNewThread);
@@ -82,7 +83,7 @@ void should_instrument_callable() throws ExecutionException, InterruptedExceptio
     }
 
     @Test
-    void should_instrument_executor() throws InterruptedException {
+    void should_instrument_executor() throws InterruptedException, TimeoutException {
         StringThreadLocalHolder.setValue("hello");
         AtomicReference<String> valueInNewThread = new AtomicReference<>();
         Executor executor = command -> new Thread(command).start();
@@ -139,10 +140,12 @@ void should_instrument_scheduled_executor_service()
         }
     }
 
-    private void runInNewThread(Runnable runnable) throws InterruptedException {
-        Thread thread = new Thread(runnable);
+    private void runInNewThread(Runnable runnable) throws InterruptedException, TimeoutException {
+        CountDownLatch latch = new CountDownLatch(1);
+        Thread thread = new Thread(countDownWhenDone(runnable, latch));
         thread.start();
-        Thread.sleep(5);
+
+        throwIfTimesOut(latch);
     }
 
     private void runInNewThread(Callable<?> callable)
@@ -157,9 +160,23 @@ private void runInNewThread(Callable<?> callable)
     }
 
     private void runInNewThread(Executor executor, AtomicReference<String> valueInNewThread)
-            throws InterruptedException {
-        executor.execute(runnable(valueInNewThread));
-        Thread.sleep(5);
+            throws InterruptedException, TimeoutException {
+        CountDownLatch latch = new CountDownLatch(1);
+        executor.execute(countDownWhenDone(runnable(valueInNewThread), latch));
+        throwIfTimesOut(latch);
+    }
+
+    private Runnable countDownWhenDone(Runnable runnable, CountDownLatch latch) {
+        return () -> {
+            runnable.run();
+            latch.countDown();
+        };
+    }
+
+    private void throwIfTimesOut(CountDownLatch latch) throws InterruptedException, TimeoutException {
+        if (!latch.await(5, TimeUnit.MILLISECONDS)) {
+            throw new TimeoutException("Waiting for executed task timed out");
+        }
     }
 
     private void runInNewThread(ExecutorService executor, AtomicReference<String> valueInNewThread,
@@ -169,8 +186,10 @@ private void runInNewThread(ExecutorService executor, AtomicReference<String> va
         StringThreadLocalHolder.setValue("hello"); // IMPORTANT: We are setting the
                                                    // thread local value as late as
                                                    // possible
-        executor.execute(runnable(valueInNewThread));
-        Thread.sleep(5);
+
+        CountDownLatch latch = new CountDownLatch(1);
+        executor.execute(countDownWhenDone(runnable(valueInNewThread), latch));
+        throwIfTimesOut(latch);
         assertion.accept(valueInNewThread);
 
         executor.submit(runnable(valueInNewThread)).get(5, TimeUnit.MILLISECONDS);