Add convenience methods to clean up database containers. Refs #159.

testcontainers · Aug 16, 2016 · f1bc5c6 · f1bc5c6
1 parent 32d252c
commit f1bc5c6
Show file tree

Hide file tree

Showing 3 changed files with 52 additions and 13 deletions.
diff --git a/modules/jdbc/src/main/java/org/testcontainers/jdbc/ContainerDatabaseDriver.java b/modules/jdbc/src/main/java/org/testcontainers/jdbc/ContainerDatabaseDriver.java
@@ -48,9 +48,9 @@ public class ContainerDatabaseDriver implements Driver {
     private static final org.slf4j.Logger LOGGER = LoggerFactory.getLogger(ContainerDatabaseDriver.class);
 
     private Driver delegate;
-    private static final Map<JdbcDatabaseContainer, Set<Connection>> containerConnections = new HashMap<>();
+    private static final Map<String, Set<Connection>> containerConnections = new HashMap<>();
     private static final Map<String, JdbcDatabaseContainer> jdbcUrlContainerCache = new HashMap<>();
-    private static final Set<JdbcDatabaseContainer> initializedContainers = new HashSet<>();
+    private static final Set<String> initializedContainers = new HashSet<>();
 
     static {
         load();
@@ -151,10 +151,10 @@ public synchronized Connection connect(String url, final Properties info) throws
               If this container has not been initialized, AND
               an init script or function has been specified, use it
              */
-            if (!initializedContainers.contains(container)) {
+            if (!initializedContainers.contains(container.getContainerId())) {
                 runInitScriptIfRequired(url, connection);
                 runInitFunctionIfRequired(url, connection);
-                initializedContainers.add(container);
+                initializedContainers.add(container.getContainerId());
             }
 
             return wrapConnection(connection, container, url);
@@ -186,11 +186,11 @@ private Map<String, String> getContainerParameters(String url) {
      * @return the connection, wrapped
      */
     private Connection wrapConnection(final Connection connection, final JdbcDatabaseContainer container, final String url) {
-        Set<Connection> connections = containerConnections.get(container);
+        Set<Connection> connections = containerConnections.get(container.getContainerId());
 
         if (connections == null) {
             connections = new HashSet<>();
-            containerConnections.put(container, connections);
+            containerConnections.put(container.getContainerId(), connections);
         }
 
         connections.add(connection);
@@ -280,4 +280,37 @@ public boolean jdbcCompliant() {
     public Logger getParentLogger() throws SQLFeatureNotSupportedException {
         return delegate.getParentLogger();
     }
+
+    /**
+     * Utility method to kill ALL database containers directly from test support code. It shouldn't be necessary to use this,
+     * but it is provided for convenience - e.g. for situations where many different database containers are being
+     * tested and cleanup is needed to limit resource usage.
+     */
+    public static void killContainers() {
+        synchronized (jdbcUrlContainerCache) {
+            jdbcUrlContainerCache.values().forEach(JdbcDatabaseContainer::stop);
+            jdbcUrlContainerCache.clear();
+            containerConnections.clear();
+            initializedContainers.clear();
+        }
+
+    }
+
+    /**
+     * Utility method to kill a database container directly from test support code. It shouldn't be necessary to use this,
+     * but it is provided for convenience - e.g. for situations where many different database containers are being
+     * tested and cleanup is needed to limit resource usage.
+     * @param jdbcUrl the JDBC URL of the container which should be killed
+     */
+    public static void killContainer(String jdbcUrl) {
+        synchronized (jdbcUrlContainerCache) {
+            JdbcDatabaseContainer container = jdbcUrlContainerCache.get(jdbcUrl);
+            if (container != null) {
+                container.stop();
+                jdbcUrlContainerCache.remove(jdbcUrl);
+                containerConnections.remove(container.getContainerId());
+                initializedContainers.remove(container.getContainerId());
+            }
+        }
+    }
 }
diff --git a/modules/jdbc/src/test/java/org/testcontainers/jdbc/ContainerDatabaseDriverTest.java b/modules/jdbc/src/test/java/org/testcontainers/jdbc/ContainerDatabaseDriverTest.java
@@ -1,16 +1,16 @@
 package org.testcontainers.jdbc;
 
-import java.sql.Connection;
-import java.sql.DriverManager;
-import java.sql.SQLException;
-import java.util.Properties;
-
 import org.hamcrest.CoreMatchers;
 import org.junit.Assert;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
 
+import java.sql.Connection;
+import java.sql.DriverManager;
+import java.sql.SQLException;
+import java.util.Properties;
+
 public class ContainerDatabaseDriverTest {
 
 	private static final String PLAIN_POSTGRESQL_JDBC_URL = "jdbc:postgresql://localhost:5432/test";

diff --git a/modules/mysql/src/test/java/org/testcontainers/jdbc/JDBCDriverTest.java b/modules/mysql/src/test/java/org/testcontainers/jdbc/JDBCDriverTest.java
@@ -5,21 +5,27 @@
 import org.apache.commons.dbutils.QueryRunner;
 import org.apache.commons.dbutils.ResultSetHandler;
 import org.apache.commons.lang.SystemUtils;
+import org.junit.After;
 import org.junit.Test;
 
 import java.sql.Connection;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
 
-import static org.rnorth.visibleassertions.VisibleAssertions.assertEquals;
 import static org.junit.Assume.assumeFalse;
+import static org.rnorth.visibleassertions.VisibleAssertions.assertEquals;
 
 /**
  *
  */
 public class JDBCDriverTest {
 
+    @After
+    public void testCleanup() {
+        ContainerDatabaseDriver.killContainers();
+    }
+
     @Test
     public void testMySQLWithVersion() throws SQLException {
         performSimpleTest("jdbc:tc:mysql:5.5.43://hostname/databasename");
@@ -32,7 +38,7 @@ public void testMySQLWithNoSpecifiedVersion() throws SQLException {
 
     @Test
     public void testMySQLWithCustomIniFile() throws SQLException {
-    	assumeFalse(SystemUtils.IS_OS_WINDOWS);
+        assumeFalse(SystemUtils.IS_OS_WINDOWS);
         HikariDataSource ds = getDataSource("jdbc:tc:mysql:5.6://hostname/databasename?TC_MY_CNF=somepath/mysql_conf_override", 1);
         Statement statement = ds.getConnection().createStatement();
         statement.execute("SELECT @@GLOBAL.innodb_file_format");