Further refine test containers

spring-boot-project/spring-boot-tools/spring-boot-test-support/src/main/java/org/springframework/boot/testsupport/testcontainers/CassandraContainer.java

@@ -0,0 +1,75 @@
+/*
+ * Copyright 2012-2018 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.boot.testsupport.testcontainers;
+
+import java.util.concurrent.Callable;
+import java.util.concurrent.TimeUnit;
+
+import com.datastax.driver.core.Cluster;
+import com.datastax.driver.core.exceptions.NoHostAvailableException;
+import org.rnorth.ducttape.TimeoutException;
+import org.rnorth.ducttape.unreliables.Unreliables;
+import org.testcontainers.containers.GenericContainer;
+import org.testcontainers.containers.wait.HostPortWaitStrategy;
+
+/**
+ * A {@link GenericContainer} for Cassandra.
+ *
+ * @author Andy Wilkinson
+ * @author Madhura Bhave
+ */
+public class CassandraContainer extends Container {
+
+	private static final int PORT = 9042;
+
+	public CassandraContainer() {
+		super("cassandra:3.11.1", PORT,
+				(container) -> container.waitingFor(new WaitStrategy()));
+	}
+
+	private static class WaitStrategy extends HostPortWaitStrategy {
+
+		@Override
+		protected void waitUntilReady() {
+			super.waitUntilReady();
+
+			try {
+				Unreliables.retryUntilTrue((int) this.startupTimeout.getSeconds(),
+						TimeUnit.SECONDS, checkConnection());
+			}
+			catch (TimeoutException ex) {
+				throw new IllegalStateException(ex);
+			}
+		}
+
+		private Callable<Boolean> checkConnection() {
+			return () -> {
+				try (Cluster cluster = Cluster.builder()
+						.withPort(this.container.getMappedPort(PORT))
+						.addContactPoint("localhost").build()) {
+					cluster.connect();
+					return true;
+				}
+				catch (IllegalArgumentException | NoHostAvailableException ex) {
+					return false;
+				}
+			};
+		}
+
+	}
+
+}

spring-boot-project/spring-boot-tools/spring-boot-test-support/src/main/java/org/springframework/boot/testsupport/testcontainers/Container.java

@@ -16,6 +16,7 @@
 
 package org.springframework.boot.testsupport.testcontainers;
 
+import java.util.function.Consumer;
 import java.util.function.Supplier;
 
 import org.junit.AssumptionViolatedException;
@@ -29,17 +30,31 @@ import org.testcontainers.containers.GenericContainer;
  * {@link TestRule} for working with an optional Docker environment. Spins up a
  * {@link GenericContainer} if a valid docker environment is found.
  *
- * @param <T> the type of the container
  * @author Madhura Bhave
+ * @author Phillip Webb
  */
-public class DockerTestContainer<T extends GenericContainer<?>> implements TestRule {
+class Container implements TestRule {
 
-	private final Supplier<T> containerSupplier;
+	private final int port;
 
-	private T container;
+	private final Supplier<GenericContainer<?>> containerFactory;
 
-	public DockerTestContainer(Supplier<T> containerSupplier) {
-		this.containerSupplier = containerSupplier;
+	private GenericContainer<?> container;
+
+	<T extends GenericContainer<T>> Container(String dockerImageName, int port) {
+		this(dockerImageName, port, null);
+	}
+
+	@SuppressWarnings({ "unchecked", "resource" })
+	<T extends GenericContainer<T>> Container(String dockerImageName, int port,
+			Consumer<T> customizer) {
+		this.port = port;
+		this.containerFactory = () -> {
+			T container = (T) new GenericContainer<>(dockerImageName)
+					.withExposedPorts(port);
+			customizer.accept(container);
+			return container;
+		};
 	}
 
 	@Override
@@ -50,12 +65,12 @@ public class DockerTestContainer<T extends GenericContainer<?>> implements TestR
 		catch (Throwable t) {
 			return new SkipStatement();
 		}
-		this.container = this.containerSupplier.get();
+		this.container = this.containerFactory.get();
 		return this.container.apply(base, description);
 	}
 
-	public int getMappedPort(int originalPort) {
-		return this.container.getMappedPort(originalPort);
+	public int getMappedPort() {
+		return this.container.getMappedPort(this.port);
 	}
 
 	private static class SkipStatement extends Statement {

spring-boot-project/spring-boot-tools/spring-boot-test-support/src/main/java/org/springframework/boot/testsupport/testcontainers/Neo4jContainer.java

@@ -19,8 +19,6 @@ package org.springframework.boot.testsupport.testcontainers;
 import java.util.concurrent.Callable;
 import java.util.concurrent.TimeUnit;
 
-import com.datastax.driver.core.Cluster;
-import com.datastax.driver.core.exceptions.NoHostAvailableException;
 import org.neo4j.ogm.config.Configuration;
 import org.neo4j.ogm.session.SessionFactory;
 import org.rnorth.ducttape.TimeoutException;
@@ -29,69 +27,26 @@ import org.testcontainers.containers.GenericContainer;
 import org.testcontainers.containers.wait.HostPortWaitStrategy;
 
 /**
- * Provides utility methods that allow creation of docker containers for
- * tests.
+ * A {@link GenericContainer} for Neo4J.
  *
  * @author Andy Wilkinson
  * @author Madhura Bhave
  */
-public abstract class TestContainers {
+public class Neo4jContainer extends Container {
 
-	@SuppressWarnings("resource")
-	public static GenericContainer<?> redis() {
-		return new GenericContainer<>("redis:4.0.6").withExposedPorts(6379);
+	public Neo4jContainer() {
+		super("neo4j:3.3.1", 7687, (container) -> container.waitingFor(new WaitStrategy())
+				.withEnv("NEO4J_AUTH", "none"));
 	}
 
-	@SuppressWarnings("resource")
-	public static GenericContainer<?> cassandra() {
-		return new GenericContainer<>("cassandra:3.11.1").withExposedPorts(9042)
-				.waitingFor(new CassandraConnectionVerifyingWaitStrategy());
-	}
-
-	public static GenericContainer<?> neo4j() {
-		return new GenericContainer<>("neo4j:3.3.1").withExposedPorts(7687)
-				.waitingFor(new Neo4jConnectionVerifyingWaitStrategy())
-				.withEnv("NEO4J_AUTH", "none");
-	}
-
-	private static class CassandraConnectionVerifyingWaitStrategy extends HostPortWaitStrategy {
-
-		@Override
-		protected void waitUntilReady() {
-			super.waitUntilReady();
-
-			try {
-				Unreliables.retryUntilTrue((int) this.startupTimeout.getSeconds(),
-						TimeUnit.SECONDS, checkConnection());
-			}
-			catch (TimeoutException ex) {
-				throw new IllegalStateException(ex);
-			}
-		}
-
-		private Callable<Boolean> checkConnection() {
-			return () -> {
-				try (Cluster cluster = Cluster.builder().withPort(this.container.getMappedPort(9042))
-						.addContactPoint("localhost")
-						.build()) {
-					cluster.connect();
-					return true;
-				}
-				catch (IllegalArgumentException | NoHostAvailableException ex) {
-					return false;
-				}
-			};
-		}
-
-	}
-
-	private static class Neo4jConnectionVerifyingWaitStrategy extends HostPortWaitStrategy {
+	private static class WaitStrategy extends HostPortWaitStrategy {
 
 		@Override
 		protected void waitUntilReady() {
 			super.waitUntilReady();
 			Configuration configuration = new Configuration.Builder()
-					.uri("bolt://localhost:" + this.container.getMappedPort(7687)).build();
+					.uri("bolt://localhost:" + this.container.getMappedPort(7687))
+					.build();
 			SessionFactory sessionFactory = new SessionFactory(configuration,
 					"org.springframework.boot.test.autoconfigure.data.neo4j");
 			try {
@@ -114,6 +69,7 @@ public abstract class TestContainers {
 				}
 			};
 		}
+
 	}
 
 }

spring-boot-project/spring-boot-tools/spring-boot-test-support/src/main/java/org/springframework/boot/testsupport/testcontainers/RedisContainer.java

@@ -0,0 +1,33 @@
+/*
+ * Copyright 2012-2018 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.boot.testsupport.testcontainers;
+
+import org.testcontainers.containers.GenericContainer;
+
+/**
+ * A {@link GenericContainer} for Redis.
+ *
+ * @author Andy Wilkinson
+ * @author Madhura Bhave
+ */
+public class RedisContainer extends Container {
+
+	public RedisContainer() {
+		super("redis:4.0.6", 6379);
+	}
+
+}