DATAMONGO-1039 - Polish db clean hook implementation.

- Refactored internal structure.
- Updated documentation.
- Added some tests

Original pull request: #222.

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/test/util/CleanMongoDB.java

@@ -18,6 +18,7 @@ package org.springframework.data.mongodb.test.util;
 import java.net.UnknownHostException;
 import java.util.Arrays;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.Set;
 
@@ -27,25 +28,35 @@ import org.junit.runners.model.Statement;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.util.CollectionUtils;
+import org.springframework.util.StringUtils;
 
 import com.mongodb.DB;
+import com.mongodb.DBCollection;
 import com.mongodb.MongoClient;
 
 /**
+ * {@link CleanMongoDB} is a junit {@link TestRule} implementation to be used as for wiping data from MongoDB instance.
+ * MongoDB specific system databases like {@literal admin} and {@literal local} remain untouched. The rule will apply
+ * <strong>after</strong> the base {@link Statement}. <br />
+ * Use as {@link org.junit.ClassRule} to wipe data after finishing all tests within a class or as {@link org.junit.Rule}
+ * to do so after each {@link org.junit.Test}.
+ * 
  * @author Christoph Strobl
+ * @since 1.6
  */
 public class CleanMongoDB implements TestRule {
 
 	private static final Logger LOGGER = LoggerFactory.getLogger(CleanMongoDB.class);
 
-	public enum Types {
+	/**
+	 * Defines contents of MongoDB.
+	 */
+	public enum Struct {
 		DATABASE, COLLECTION, INDEX;
 	}
 
+	@SuppressWarnings("serial")//
 	private Set<String> preserveDatabases = new HashSet<String>() {
-
-		private static final long serialVersionUID = -8698807376808700046L;
-
 		{
 			add("admin");
 			add("local");
@@ -54,57 +65,278 @@ public class CleanMongoDB implements TestRule {
 
 	private Set<String> dbNames = new HashSet<String>();
 	private Set<String> collectionNames = new HashSet<String>();
-	private Set<Types> types = new HashSet<CleanMongoDB.Types>();
+	private Set<Struct> types = new HashSet<CleanMongoDB.Struct>();
 	private MongoClient client;
 
+	/**
+	 * Create new instance using an internal {@link MongoClient}.
+	 */
 	public CleanMongoDB() {
 		this(null);
 	}
 
+	/**
+	 * Create new instance using an internal {@link MongoClient} connecting to specified instance running at host:port.
+	 * 
+	 * @param host
+	 * @param port
+	 * @throws UnknownHostException
+	 */
 	public CleanMongoDB(String host, int port) throws UnknownHostException {
 		this(new MongoClient(host, port));
 	}
 
+	/**
+	 * Create new instance using the given client.
+	 * 
+	 * @param client
+	 */
 	public CleanMongoDB(MongoClient client) {
 		this.client = client;
 	}
 
+	/**
+	 * Removes everything by dropping every single {@link DB}.
+	 * 
+	 * @return
+	 */
 	public static CleanMongoDB everything() {
 
 		CleanMongoDB cleanMongoDB = new CleanMongoDB();
-		cleanMongoDB.clean(Types.DATABASE);
+		cleanMongoDB.clean(Struct.DATABASE);
 		return cleanMongoDB;
 	}
 
+	/**
+	 * Removes everything from the databases with given name by dropping the according {@link DB}.
+	 * 
+	 * @param dbNames
+	 * @return
+	 */
 	public static CleanMongoDB databases(String... dbNames) {
 
 		CleanMongoDB cleanMongoDB = new CleanMongoDB();
-		cleanMongoDB.clean(Types.DATABASE);
+		cleanMongoDB.clean(Struct.DATABASE);
-		cleanMongoDB.collectionNames.addAll(Arrays.asList(dbNames));
+		cleanMongoDB.useDatabases(dbNames);
+		return cleanMongoDB;
+	}
+
+	/**
+	 * Drops the {@link DBCollection} with given names from every single {@link DB} containing them.
+	 * 
+	 * @param collectionNames
+	 * @return
+	 */
+	public static CleanMongoDB collections(String... collectionNames) {
+		return collections("", Arrays.asList(collectionNames));
+	}
+
+	/**
+	 * Drops the {@link DBCollection} with given names from the named {@link DB}.
+	 * 
+	 * @param dbName
+	 * @param collectionNames
+	 * @return
+	 */
+	public static CleanMongoDB collections(String dbName, Collection<String> collectionNames) {
+
+		CleanMongoDB cleanMongoDB = new CleanMongoDB();
+		cleanMongoDB.clean(Struct.COLLECTION);
+		cleanMongoDB.useCollections(dbName, collectionNames);
 		return cleanMongoDB;
 	}
 
+	/**
+	 * Drops all index structures from every single {@link DBCollection}.
+	 * 
+	 * @return
+	 */
 	public static CleanMongoDB indexes() {
+		return indexes(Collections.<String> emptySet());
+	}
+
+	/**
+	 * Drops all index structures from every single {@link DBCollection}.
+	 * 
+	 * @param collectionNames
+	 * @return
+	 */
+	public static CleanMongoDB indexes(Collection<String> collectionNames) {
 
 		CleanMongoDB cleanMongoDB = new CleanMongoDB();
-		cleanMongoDB.clean(Types.INDEX);
+		cleanMongoDB.clean(Struct.INDEX);
+		cleanMongoDB.useCollections(collectionNames);
 		return cleanMongoDB;
 	}
 
-	public CleanMongoDB clean(Types... types) {
+	/**
+	 * Define {@link Struct} to be cleaned.
+	 * 
+	 * @param types
+	 * @return
+	 */
+	public CleanMongoDB clean(Struct... types) {
 
 		this.types.addAll(Arrays.asList(types));
 		return this;
 	}
 
-	public Statement apply() {
+	/**
+	 * Defines the {@link DB}s to be used. <br />
+	 * Impact along with {@link CleanMongoDB#clean(Struct...)}:
+	 * <ul>
+	 * <li>{@link Struct#DATABASE}: Forces drop of named databases.</li>
+	 * <li>{@link Struct#COLLECTION}: Forces drop of collections within named databases.</li>
+	 * <li>{@link Struct#INDEX}: Removes index within collections of named databases.</li>
+	 * </ul>
+	 * 
+	 * @param dbNames
+	 * @return
+	 */
+	public CleanMongoDB useDatabases(String... dbNames) {
+
+		this.dbNames.addAll(Arrays.asList(dbNames));
+		return this;
+	}
+
+	/**
+	 * Excludes the given {@link DB}s from being processed.
+	 * 
+	 * @param dbNames
+	 * @return
+	 */
+	public CleanMongoDB preserveDatabases(String... dbNames) {
+		this.preserveDatabases.addAll(Arrays.asList(dbNames));
+		return this;
+	}
+
+	/**
+	 * Defines the {@link DBCollection}s to be used. <br />
+	 * Impact along with {@link CleanMongoDB#clean(Struct...)}:
+	 * <ul>
+	 * <li>{@link Struct#COLLECTION}: Forces drop of named collections.</li>
+	 * <li>{@link Struct#INDEX}: Removes index within named collections.</li>
+	 * </ul>
+	 * 
+	 * @param collectionNames
+	 * @return
+	 */
+	public CleanMongoDB useCollections(String... collectionNames) {
+		return useCollections(Arrays.asList(collectionNames));
+	}
+
+	private CleanMongoDB useCollections(Collection<String> collectionNames) {
+		return useCollections("", collectionNames);
+	}
+
+	/**
+	 * Defines the {@link DBCollection}s and {@link DB} to be used. <br />
+	 * Impact along with {@link CleanMongoDB#clean(Struct...)}:
+	 * <ul>
+	 * <li>{@link Struct#COLLECTION}: Forces drop of named collections in given db.</li>
+	 * <li>{@link Struct#INDEX}: Removes index within named collections in given db.</li>
+	 * </ul>
+	 * 
+	 * @param collectionNames
+	 * @return
+	 */
+	public CleanMongoDB useCollections(String db, Collection<String> collectionNames) {
+
+		if (StringUtils.hasText(db)) {
+			this.dbNames.add(db);
+		}
+
+		if (!CollectionUtils.isEmpty(collectionNames)) {
+			this.collectionNames.addAll(collectionNames);
+		}
+		return this;
+	}
+
+	Statement apply() {
 		return apply(null, null);
 	}
 
+	/*
+	 * (non-Javadoc)
+	 * @see org.junit.rules.TestRule#apply(org.junit.runners.model.Statement, org.junit.runner.Description)
+	 */
 	public Statement apply(Statement base, Description description) {
 		return new MongoCleanStatement(base);
 	}
 
+	private void doClean() {
+
+		Collection<String> dbNamesToUse = initDbNames();
+
+		for (String dbName : dbNamesToUse) {
+
+			if (isPreserved(dbName) || dropDbIfRequired(dbName)) {
+				continue;
+			}
+
+			DB db = client.getDB(dbName);
+			dropCollectionsOrIndexIfRequried(db, initCollectionNames(db));
+		}
+	}
+
+	private boolean dropDbIfRequired(String dbName) {
+
+		if (!types.contains(Struct.DATABASE)) {
+			return false;
+		}
+
+		client.dropDatabase(dbName);
+		LOGGER.debug("Dropping DB '{}'. ", dbName);
+		return true;
+	}
+
+	private void dropCollectionsOrIndexIfRequried(DB db, Collection<String> collectionsToUse) {
+
+		for (String collectionName : collectionsToUse) {
+
+			if (db.collectionExists(collectionName)) {
+
+				DBCollection collection = db.getCollectionFromString(collectionName);
+				if (collection != null) {
+
+					if (types.contains(Struct.COLLECTION)) {
+						collection.drop();
+						LOGGER.debug("Dropping collection '{}' for DB '{}'. ", collectionName, db.getName());
+					} else if (types.contains(Struct.INDEX)) {
+						collection.dropIndexes();
+						LOGGER.debug("Dropping indexes in collection '{}' for DB '{}'. ", collectionName, db.getName());
+					}
+				}
+			}
+		}
+	}
+
+	private boolean isPreserved(String dbName) {
+		return preserveDatabases.contains(dbName.toLowerCase());
+	}
+
+	private Collection<String> initDbNames() {
+
+		Collection<String> dbNamesToUse = dbNames;
+		if (dbNamesToUse.isEmpty()) {
+			dbNamesToUse = client.getDatabaseNames();
+		}
+		return dbNamesToUse;
+	}
+
+	private Collection<String> initCollectionNames(DB db) {
+
+		Collection<String> collectionsToUse = collectionNames;
+		if (CollectionUtils.isEmpty(collectionsToUse)) {
+			collectionsToUse = db.getCollectionNames();
+		}
+		return collectionsToUse;
+	}
+
+	/**
+	 * @author Christoph Strobl
+	 * @since 1.6
+	 */
 	private class MongoCleanStatement extends Statement {
 
 		private final Statement base;
@@ -126,61 +358,12 @@ public class CleanMongoDB implements TestRule {
 				isInternal = true;
 			}
 
-			Collection<String> dbNamesToUse = dbNames;
+			doClean();
-			if (dbNamesToUse.isEmpty()) {
-				dbNamesToUse = client.getDatabaseNames();
-			}
-
-			for (String dbName : dbNamesToUse) {
-
-				if (preserveDatabases.contains(dbName.toLowerCase())) {
-					continue;
-				}
-
-				if (types.contains(Types.DATABASE)) {
-					client.dropDatabase(dbName);
-					LOGGER.debug("Dropping DB '{}'. ", dbName);
-				}
-
-				if (types.contains(Types.COLLECTION)) {
-
-					DB db = client.getDB(dbName);
-					Collection<String> collectionsToUse = initCollectionNames(db);
-					for (String collectionName : collectionsToUse) {
-						if (db.collectionExists(collectionName)) {
-							db.getCollectionFromString(collectionName).drop();
-							LOGGER.debug("Dropping collection '{}' for DB '{}'. ", collectionName, dbName);
-						}
-					}
-				}
-
-				if (types.contains(Types.INDEX)) {
-
-					DB db = client.getDB(dbName);
-					Collection<String> collectionsToUse = initCollectionNames(db);
-					for (String collectionName : collectionsToUse) {
-						if (db.collectionExists(collectionName)) {
-							db.getCollectionFromString(collectionName).dropIndexes();
-							LOGGER.debug("Dropping indexes in collection '{}' for DB '{}'. ", collectionName, dbName);
-						}
-					}
-				}
-			}
 
 			if (isInternal) {
 				client.close();
 				client = null;
 			}
 		}
-
-		private Collection<String> initCollectionNames(DB db) {
-
-			Collection<String> collectionsToUse = collectionNames;
-			if (CollectionUtils.isEmpty(collectionsToUse)) {
-				collectionsToUse = db.getCollectionNames();
-			}
-			return collectionsToUse;
-		}
 	}
-
 }

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/test/util/CleanMongoDBJunitRunListener.java

@@ -17,21 +17,24 @@ package org.springframework.data.mongodb.test.util;
 
 import org.junit.runner.Result;
 import org.junit.runner.notification.RunListener;
-import org.springframework.data.mongodb.test.util.CleanMongoDB.Types;
+import org.springframework.data.mongodb.test.util.CleanMongoDB.Struct;
 
 /**
+ * {@link RunListener} implementation to be used for wiping MongoDB index structures after all test runs have finished.
+ * 
  * @author Christoph Strobl
+ * @since 1.6
  */
 public class CleanMongoDBJunitRunListener extends RunListener {
 
 	@Override
 	public void testRunFinished(Result result) throws Exception {
+
 		super.testRunFinished(result);
 		try {
-			new CleanMongoDB().clean(Types.INDEX).apply().evaluate();
+			new CleanMongoDB().clean(Struct.INDEX).apply().evaluate();
 		} catch (Throwable e) {
 			e.printStackTrace();
 		}
 	}
-
 }

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/test/util/CleanMongoDBTests.java

@@ -20,6 +20,7 @@ import static org.mockito.Mockito.*;
 
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashSet;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -28,7 +29,7 @@ import org.junit.runner.RunWith;
 import org.junit.runners.model.Statement;
 import org.mockito.Mock;
 import org.mockito.runners.MockitoJUnitRunner;
-import org.springframework.data.mongodb.test.util.CleanMongoDB.Types;
+import org.springframework.data.mongodb.test.util.CleanMongoDB.Struct;
 
 import com.mongodb.DB;
 import com.mongodb.DBCollection;
@@ -41,58 +42,146 @@ import com.mongodb.MongoClient;
 public class CleanMongoDBTests {
 
 	private CleanMongoDB cleaner;
+
+	// JUnit internals
 	private @Mock Statement baseStatementMock;
 	private @Mock Description descriptionMock;
+
+	// MongoClient in use
 	private @Mock MongoClient mongoClientMock;
-	private @Mock DB db1mock;
-	private @Mock DB db2mock;
-	private @Mock DBCollection collection1mock;
 
+	// Some Mock DBs
+	private @Mock DB db1mock, db2mock;
+	private @Mock DBCollection db1collection1mock, db1collection2mock, db2collection1mock;
+
+	@SuppressWarnings("serial")
 	@Before
 	public void setUp() {
 
+		// DB setup
 		when(mongoClientMock.getDatabaseNames()).thenReturn(Arrays.asList("admin", "db1", "db2"));
 		when(mongoClientMock.getDB(eq("db1"))).thenReturn(db1mock);
 		when(mongoClientMock.getDB(eq("db2"))).thenReturn(db2mock);
+
+		// collections have to exist
 		when(db1mock.collectionExists(anyString())).thenReturn(true);
 		when(db2mock.collectionExists(anyString())).thenReturn(true);
-		when(db1mock.getCollectionNames()).thenReturn(Collections.singleton("collection-1"));
+
-		when(db2mock.getCollectionNames()).thenReturn(Collections.<String> emptySet());
+		// init collection names per database
-		when(db1mock.getCollectionFromString(eq("collection-1"))).thenReturn(collection1mock);
+		when(db1mock.getCollectionNames()).thenReturn(new HashSet<String>() {
+			{
+				add("db1collection1");
+				add("db1collection2");
+			}
+		});
+		when(db2mock.getCollectionNames()).thenReturn(Collections.singleton("db2collection1"));
+
+		// return collections according to names
+		when(db1mock.getCollectionFromString(eq("db1collection1"))).thenReturn(db1collection1mock);
+		when(db1mock.getCollectionFromString(eq("db1collection2"))).thenReturn(db1collection2mock);
+		when(db2mock.getCollectionFromString(eq("db2collection1"))).thenReturn(db2collection1mock);
 
 		cleaner = new CleanMongoDB(mongoClientMock);
 	}
 
 	@Test
-	public void preservesSystemCollectionsCorrectly() throws Throwable {
+	public void preservesSystemDBsCorrectlyWhenCleaningDatabase() throws Throwable {
+
+		cleaner.clean(Struct.DATABASE);
+
+		cleaner.apply(baseStatementMock, descriptionMock).evaluate();
+
+		verify(mongoClientMock, never()).dropDatabase(eq("admin"));
+	}
+
+	@Test
+	public void preservesNamedDBsCorrectlyWhenCleaningDatabase() throws Throwable {
 
-		cleaner.clean(Types.DATABASE);
+		cleaner.clean(Struct.DATABASE);
+		cleaner.preserveDatabases("db1");
+
+		cleaner.apply(baseStatementMock, descriptionMock).evaluate();
+
+		verify(mongoClientMock, never()).dropDatabase(eq("db1"));
+	}
+
+	@Test
+	public void dropsAllDBsCorrectlyWhenCleaingDatabaseAndNotExplictDBNamePresent() throws Throwable {
+
+		cleaner.clean(Struct.DATABASE);
 
 		cleaner.apply(baseStatementMock, descriptionMock).evaluate();
 
 		verify(mongoClientMock, times(1)).dropDatabase(eq("db1"));
 		verify(mongoClientMock, times(1)).dropDatabase(eq("db2"));
-		verify(mongoClientMock, never()).dropDatabase(eq("admin"));
 	}
 
 	@Test
-	public void removesCollectionsCorrectly() throws Throwable {
+	public void dropsSpecifiedDBsCorrectlyWhenExplicitNameSet() throws Throwable {
+
+		cleaner.clean(Struct.DATABASE);
+		cleaner.useDatabases("db2");
+
+		cleaner.apply(baseStatementMock, descriptionMock).evaluate();
+
+		verify(mongoClientMock, times(1)).dropDatabase(eq("db2"));
+		verify(mongoClientMock, never()).dropDatabase(eq("db1"));
+	}
+
+	@Test
+	public void doesNotRemoveAnyDBwhenCleaningCollections() throws Throwable {
 
-		cleaner.clean(Types.COLLECTION);
+		cleaner.clean(Struct.COLLECTION);
 
 		cleaner.apply(baseStatementMock, descriptionMock).evaluate();
 
 		verify(mongoClientMock, never()).dropDatabase(eq("db1"));
 		verify(mongoClientMock, never()).dropDatabase(eq("db2"));
 		verify(mongoClientMock, never()).dropDatabase(eq("admin"));
+	}
+
+	@Test
+	public void doesNotDropCollectionsFromPreservedDBs() throws Throwable {
+
+		cleaner.clean(Struct.COLLECTION);
+		cleaner.preserveDatabases("db1");
+
+		cleaner.apply(baseStatementMock, descriptionMock).evaluate();
+
+		verify(db1collection1mock, never()).drop();
+		verify(db1collection2mock, never()).drop();
+		verify(db2collection1mock, times(1)).drop();
+	}
+
+	@Test
+	public void removesAllCollectionsFromAllDatabasesWhenNotLimitedToSpecificOnes() throws Throwable {
+
+		cleaner.clean(Struct.COLLECTION);
+
+		cleaner.apply(baseStatementMock, descriptionMock).evaluate();
+
+		verify(db1collection1mock, times(1)).drop();
+		verify(db1collection2mock, times(1)).drop();
+		verify(db2collection1mock, times(1)).drop();
+	}
+
+	@Test
+	public void removesOnlyNamedCollectionsWhenSpecified() throws Throwable {
+
+		cleaner.clean(Struct.COLLECTION);
+		cleaner.useCollections("db1collection2");
+
+		cleaner.apply(baseStatementMock, descriptionMock).evaluate();
 
-		verify(collection1mock, times(1)).drop();
+		verify(db1collection1mock, never()).drop();
+		verify(db2collection1mock, never()).drop();
+		verify(db1collection2mock, times(1)).drop();
 	}
 
 	@Test
 	public void removesIndexesCorrectly() throws Throwable {
 
-		cleaner.clean(Types.INDEX);
+		cleaner.clean(Struct.INDEX);
 
 		cleaner.apply(baseStatementMock, descriptionMock).evaluate();
 
@@ -100,6 +189,6 @@ public class CleanMongoDBTests {
 		verify(mongoClientMock, never()).dropDatabase(eq("db2"));
 		verify(mongoClientMock, never()).dropDatabase(eq("admin"));
 
-		verify(collection1mock, times(1)).dropIndexes();
+		verify(db1collection1mock, times(1)).dropIndexes();
 	}
 }