From b291332cd41fd747b6665b4916c460a35ac7f851 Mon Sep 17 00:00:00 2001
From: Dave Syer <dsyer@gopivotal.com>
Date: Wed, 23 Apr 2014 10:13:17 +0100
Subject: [PATCH] Use CountDownLatch instead of Thread.sleep()

... to wait for ApplicationContext to close in the 3 tests
that we needed to do so.

Fixes gh-664
---
 .../endpoint/ShutdownEndpointTests.java       | 23 +++++++++++++---
 .../endpoint/ShutdownParentEndpointTests.java | 26 +++++++++++++++----
 2 files changed, 41 insertions(+), 8 deletions(-)

diff --git a/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/endpoint/ShutdownEndpointTests.java b/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/endpoint/ShutdownEndpointTests.java
index da2f64f8751..f3eede3bebc 100644
--- a/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/endpoint/ShutdownEndpointTests.java
+++ b/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/endpoint/ShutdownEndpointTests.java
@@ -16,13 +16,17 @@
 
 package org.springframework.boot.actuate.endpoint;
 
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+
 import org.junit.Test;
 import org.springframework.boot.context.properties.EnableConfigurationProperties;
+import org.springframework.context.ApplicationListener;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
+import org.springframework.context.event.ContextClosedEvent;
 
 import static org.hamcrest.Matchers.startsWith;
-import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertThat;
 import static org.junit.Assert.assertTrue;
 
@@ -41,17 +45,19 @@ public class ShutdownEndpointTests extends AbstractEndpointTests<ShutdownEndpoin
 
 	@Test
 	public void invoke() throws Exception {
+		CountDownLatch latch = this.context.getBean(Config.class).latch;
 		assertThat((String) getEndpointBean().invoke().get("message"),
 				startsWith("Shutting down"));
 		assertTrue(this.context.isActive());
-		Thread.sleep(600);
-		assertFalse(this.context.isActive());
+		assertTrue(latch.await(10, TimeUnit.SECONDS));
 	}
 
 	@Configuration
 	@EnableConfigurationProperties
 	public static class Config {
 
+		private CountDownLatch latch = new CountDownLatch(1);
+
 		@Bean
 		public ShutdownEndpoint endpoint() {
 			ShutdownEndpoint endpoint = new ShutdownEndpoint();
@@ -59,5 +65,16 @@ public class ShutdownEndpointTests extends AbstractEndpointTests<ShutdownEndpoin
 			return endpoint;
 		}
 
+		@Bean
+		public ApplicationListener<ContextClosedEvent> listener() {
+			return new ApplicationListener<ContextClosedEvent>() {
+				@Override
+				public void onApplicationEvent(ContextClosedEvent event) {
+					Config.this.latch.countDown();
+				}
+			};
+
+		}
+
 	}
 }
diff --git a/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/endpoint/ShutdownParentEndpointTests.java b/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/endpoint/ShutdownParentEndpointTests.java
index c782dfe3352..215cb1aa279 100644
--- a/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/endpoint/ShutdownParentEndpointTests.java
+++ b/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/endpoint/ShutdownParentEndpointTests.java
@@ -16,16 +16,20 @@
 
 package org.springframework.boot.actuate.endpoint;
 
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+
 import org.junit.After;
 import org.junit.Test;
 import org.springframework.boot.builder.SpringApplicationBuilder;
 import org.springframework.boot.context.properties.EnableConfigurationProperties;
+import org.springframework.context.ApplicationListener;
 import org.springframework.context.ConfigurableApplicationContext;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
+import org.springframework.context.event.ContextClosedEvent;
 
 import static org.hamcrest.Matchers.startsWith;
-import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertThat;
 import static org.junit.Assert.assertTrue;
 
@@ -49,22 +53,22 @@ public class ShutdownParentEndpointTests {
 	public void shutdownChild() throws Exception {
 		this.context = new SpringApplicationBuilder(Config.class).child(Empty.class)
 				.web(false).run();
+		CountDownLatch latch = this.context.getBean(Config.class).latch;
 		assertThat((String) getEndpointBean().invoke().get("message"),
 				startsWith("Shutting down"));
 		assertTrue(this.context.isActive());
-		Thread.sleep(600);
-		assertFalse(this.context.isActive());
+		assertTrue(latch.await(10, TimeUnit.SECONDS));
 	}
 
 	@Test
 	public void shutdownParent() throws Exception {
 		this.context = new SpringApplicationBuilder(Empty.class).child(Config.class)
 				.web(false).run();
+		CountDownLatch latch = this.context.getBean(Config.class).latch;
 		assertThat((String) getEndpointBean().invoke().get("message"),
 				startsWith("Shutting down"));
 		assertTrue(this.context.isActive());
-		Thread.sleep(600);
-		assertFalse(this.context.isActive());
+		assertTrue(latch.await(10, TimeUnit.SECONDS));
 	}
 
 	private ShutdownEndpoint getEndpointBean() {
@@ -75,6 +79,8 @@ public class ShutdownParentEndpointTests {
 	@EnableConfigurationProperties
 	public static class Config {
 
+		private CountDownLatch latch = new CountDownLatch(1);
+
 		@Bean
 		public ShutdownEndpoint endpoint() {
 			ShutdownEndpoint endpoint = new ShutdownEndpoint();
@@ -82,6 +88,16 @@ public class ShutdownParentEndpointTests {
 			return endpoint;
 		}
 
+		@Bean
+		public ApplicationListener<ContextClosedEvent> listener() {
+			return new ApplicationListener<ContextClosedEvent>() {
+				@Override
+				public void onApplicationEvent(ContextClosedEvent event) {
+					Config.this.latch.countDown();
+				}
+			};
+
+		}
 	}
 
 	@Configuration