From c9f04c397793b661569b012a647926b82d5aa11b Mon Sep 17 00:00:00 2001
From: Andy Wilkinson <awilkinson@pivotal.io>
Date: Tue, 24 Apr 2018 11:03:29 +0100
Subject: [PATCH] Avoid race between container starting and getting mapped port

---
 .../testcontainers/CassandraContainer.java       | 16 +++++++++-------
 .../testcontainers/Neo4jContainer.java           | 16 ++++++++--------
 2 files changed, 17 insertions(+), 15 deletions(-)

diff --git a/spring-boot-project/spring-boot-tools/spring-boot-test-support/src/main/java/org/springframework/boot/testsupport/testcontainers/CassandraContainer.java b/spring-boot-project/spring-boot-tools/spring-boot-test-support/src/main/java/org/springframework/boot/testsupport/testcontainers/CassandraContainer.java
index 63ac1c4ba99..f6e830da791 100644
--- a/spring-boot-project/spring-boot-tools/spring-boot-test-support/src/main/java/org/springframework/boot/testsupport/testcontainers/CassandraContainer.java
+++ b/spring-boot-project/spring-boot-tools/spring-boot-test-support/src/main/java/org/springframework/boot/testsupport/testcontainers/CassandraContainer.java
@@ -38,17 +38,18 @@ public class CassandraContainer extends Container {
 	private static final int PORT = 9042;
 
 	public CassandraContainer() {
-		super("cassandra:3.11.1", PORT, (container) -> container
-				.waitingFor(new WaitStrategy(container.getMappedPort(PORT)))
-				.withStartupAttempts(3).withStartupTimeout(Duration.ofSeconds(60)));
+		super("cassandra:3.11.1", PORT,
+				(container) -> container.waitingFor(new WaitStrategy(container))
+						.withStartupAttempts(3)
+						.withStartupTimeout(Duration.ofSeconds(60)));
 	}
 
 	private static final class WaitStrategy extends HostPortWaitStrategy {
 
-		private final int port;
+		private final GenericContainer<?> container;
 
-		private WaitStrategy(int port) {
-			this.port = port;
+		private WaitStrategy(GenericContainer<?> container) {
+			this.container = container;
 		}
 
 		@Override
@@ -66,7 +67,8 @@ public class CassandraContainer extends Container {
 
 		private Callable<Boolean> checkConnection() {
 			return () -> {
-				try (Cluster cluster = Cluster.builder().withPort(this.port)
+				try (Cluster cluster = Cluster.builder()
+						.withPort(this.container.getMappedPort(CassandraContainer.PORT))
 						.addContactPoint("localhost").build()) {
 					cluster.connect();
 					return true;
diff --git a/spring-boot-project/spring-boot-tools/spring-boot-test-support/src/main/java/org/springframework/boot/testsupport/testcontainers/Neo4jContainer.java b/spring-boot-project/spring-boot-tools/spring-boot-test-support/src/main/java/org/springframework/boot/testsupport/testcontainers/Neo4jContainer.java
index 83a11943ee0..86c3d2ade09 100644
--- a/spring-boot-project/spring-boot-tools/spring-boot-test-support/src/main/java/org/springframework/boot/testsupport/testcontainers/Neo4jContainer.java
+++ b/spring-boot-project/spring-boot-tools/spring-boot-test-support/src/main/java/org/springframework/boot/testsupport/testcontainers/Neo4jContainer.java
@@ -37,25 +37,25 @@ public class Neo4jContainer extends Container {
 	private static final int PORT = 7687;
 
 	public Neo4jContainer() {
-		super("neo4j:3.3.1", PORT,
-				(container) -> container
-						.waitingFor(new WaitStrategy(container.getMappedPort(PORT)))
-						.withEnv("NEO4J_AUTH", "none"));
+		super("neo4j:3.3.1", PORT, (container) -> container
+				.waitingFor(new WaitStrategy(container)).withEnv("NEO4J_AUTH", "none"));
 	}
 
 	private static final class WaitStrategy extends HostPortWaitStrategy {
 
-		private final int port;
+		private final GenericContainer<?> container;
 
-		private WaitStrategy(int port) {
-			this.port = port;
+		private WaitStrategy(GenericContainer<?> container) {
+			this.container = container;
 		}
 
 		@Override
 		public void waitUntilReady() {
 			super.waitUntilReady();
 			Configuration configuration = new Configuration.Builder()
-					.uri("bolt://localhost:" + this.port).build();
+					.uri("bolt://localhost:"
+							+ this.container.getMappedPort(Neo4jContainer.PORT))
+					.build();
 			SessionFactory sessionFactory = new SessionFactory(configuration,
 					"org.springframework.boot.test.autoconfigure.data.neo4j");
 			try {