Java tutorial
/* * Copyright 2014 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.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; import org.junit.rules.TestRule; import org.junit.runner.Description; 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); /** * Defines contents of MongoDB. */ public enum Struct { DATABASE, COLLECTION, INDEX; } @SuppressWarnings("serial") // private Set<String> preserveDatabases = new HashSet<String>() { { add("admin"); add("local"); } }; private Set<String> dbNames = new HashSet<String>(); private Set<String> collectionNames = new HashSet<String>(); 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(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(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(Struct.INDEX); cleanMongoDB.useCollections(collectionNames); return cleanMongoDB; } /** * Define {@link Struct} to be cleaned. * * @param types * @return */ public CleanMongoDB clean(Struct... types) { this.types.addAll(Arrays.asList(types)); return this; } /** * 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; public MongoCleanStatement(Statement base) { this.base = base; } @Override public void evaluate() throws Throwable { if (base != null) { base.evaluate(); } boolean isInternal = false; if (client == null) { client = new MongoClient(); isInternal = true; } doClean(); if (isInternal) { client.close(); client = null; } } } }