DATAMONGO-1039 - Polish db clean hook implementation.

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

Original pull request: #222.

Author: christophstrobl

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

@@ -18,6 +18,7 @@
 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.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 enum Types {
 
 	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.collectionNames.addAll(Arrays.asList(dbNames));
+		cleanMongoDB.clean(Struct.DATABASE);
+		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 void evaluate() throws Throwable {
 				isInternal = true;
 			}
 
-			Collection<String> dbNamesToUse = dbNames;
-			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);
-						}
-					}
-				}
-			}
+			doClean();
 
 			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;
-		}
 	}
-
 }