path: root/contrib/storage-mongo/src/test/java/org/apache/drill/exec/store/mongo/MongoTestSuit.java

/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.apache.drill.exec.store.mongo;

import java.io.IOException;
import java.net.UnknownHostException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;

import org.bson.Document;
import org.bson.conversions.Bson;
import org.junit.AfterClass;
import org.junit.BeforeClass;
import org.junit.runner.RunWith;
import org.junit.runners.Suite;
import org.junit.runners.Suite.SuiteClasses;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.mongodb.MongoClient;
import com.mongodb.ServerAddress;
import com.mongodb.client.MongoCollection;
import com.mongodb.client.MongoDatabase;
import com.mongodb.client.model.IndexOptions;
import com.mongodb.client.model.Indexes;
import de.flapdoodle.embed.mongo.Command;
import de.flapdoodle.embed.mongo.MongodExecutable;
import de.flapdoodle.embed.mongo.MongodProcess;
import de.flapdoodle.embed.mongo.MongodStarter;
import de.flapdoodle.embed.mongo.config.IMongoCmdOptions;
import de.flapdoodle.embed.mongo.config.IMongodConfig;
import de.flapdoodle.embed.mongo.config.IMongosConfig;
import de.flapdoodle.embed.mongo.config.MongoCmdOptionsBuilder;
import de.flapdoodle.embed.mongo.config.MongodConfigBuilder;
import de.flapdoodle.embed.mongo.config.MongosConfigBuilder;
import de.flapdoodle.embed.mongo.config.Net;
import de.flapdoodle.embed.mongo.config.RuntimeConfigBuilder;
import de.flapdoodle.embed.mongo.config.Storage;
import de.flapdoodle.embed.mongo.distribution.Version;
import de.flapdoodle.embed.mongo.tests.MongosSystemForTestFactory;
import de.flapdoodle.embed.process.config.IRuntimeConfig;
import de.flapdoodle.embed.process.runtime.Network;

@RunWith(Suite.class)
@SuiteClasses({ TestMongoFilterPushDown.class, TestMongoProjectPushDown.class,
    TestMongoQueries.class, TestMongoChunkAssignment.class })
public class MongoTestSuit implements MongoTestConstants {

  private static final Logger logger = LoggerFactory
      .getLogger(MongoTestSuit.class);

  protected static MongoClient mongoClient;

  private static boolean distMode = System.getProperty(
      "drill.mongo.tests.shardMode", "true").equalsIgnoreCase("true");
  private static boolean authEnabled = System.getProperty(
      "drill.mongo.tests.authEnabled", "false").equalsIgnoreCase("true");

  private static volatile AtomicInteger initCount = new AtomicInteger(0);

  private static class DistributedMode {
    private static MongosSystemForTestFactory mongosTestFactory;

    private static void setup() throws Exception {
      // creating configServers
      List<IMongodConfig> configServers = new ArrayList<IMongodConfig>(1);
      IMongodConfig configIMongodConfig = crateConfigServerConfig(
          CONFIG_SERVER_PORT, true);
      configServers.add(configIMongodConfig);

      // creating replicaSets
      Map<String, List<IMongodConfig>> replicaSets = new HashMap<String, List<IMongodConfig>>();
      List<IMongodConfig> replicaSet1 = new ArrayList<IMongodConfig>();
      replicaSet1.add(crateIMongodConfig(MONGOD_1_PORT, false,
          REPLICA_SET_1_NAME));
      replicaSet1.add(crateIMongodConfig(MONGOD_2_PORT, false,
          REPLICA_SET_1_NAME));
      replicaSet1.add(crateIMongodConfig(MONGOD_3_PORT, false,
          REPLICA_SET_1_NAME));
      replicaSets.put(REPLICA_SET_1_NAME, replicaSet1);
      List<IMongodConfig> replicaSet2 = new ArrayList<IMongodConfig>();
      replicaSet2.add(crateIMongodConfig(MONGOD_4_PORT, false,
          REPLICA_SET_2_NAME));
      replicaSet2.add(crateIMongodConfig(MONGOD_5_PORT, false,
          REPLICA_SET_2_NAME));
      replicaSet2.add(crateIMongodConfig(MONGOD_6_PORT, false,
          REPLICA_SET_2_NAME));
      replicaSets.put(REPLICA_SET_2_NAME, replicaSet2);

      // create mongos
      IMongosConfig mongosConfig = createIMongosConfig();
      mongosTestFactory = new MongosSystemForTestFactory(mongosConfig,
          replicaSets, configServers, EMPLOYEE_DB, EMPINFO_COLLECTION,
          "employee_id");
      try {
        mongosTestFactory.start();
        mongoClient = (MongoClient) mongosTestFactory.getMongo();
      } catch (Throwable e) {
        logger.error(" Error while starting shrded cluster. ", e);
        throw new Exception(" Error while starting shrded cluster. ", e);
      }
      createDbAndCollections(DONUTS_DB, DONUTS_COLLECTION, "id");
      createDbAndCollections(EMPLOYEE_DB, EMPTY_COLLECTION, "field_2");
      createDbAndCollections(DATATYPE_DB, DATATYPE_COLLECTION, "_id");
    }

    private static IMongodConfig crateConfigServerConfig(int configServerPort,
        boolean flag) throws UnknownHostException, IOException {
      IMongoCmdOptions cmdOptions = new MongoCmdOptionsBuilder().useNoJournal(false).verbose(false)
          .build();

      IMongodConfig mongodConfig = new MongodConfigBuilder()
          .version(Version.Main.PRODUCTION)
          .net(new Net(LOCALHOST, configServerPort, Network.localhostIsIPv6()))
          .configServer(flag).cmdOptions(cmdOptions).build();
      return mongodConfig;
    }

    private static IMongodConfig crateIMongodConfig(int mongodPort,
        boolean flag, String replicaName) throws UnknownHostException,
        IOException {
      IMongoCmdOptions cmdOptions = new MongoCmdOptionsBuilder().useNoJournal(false).verbose(false)
          .build();

      Storage replication = new Storage(null, replicaName, 0);
      IMongodConfig mongodConfig = new MongodConfigBuilder()
          .version(Version.Main.PRODUCTION)
          .net(new Net(LOCALHOST, mongodPort, Network.localhostIsIPv6()))
          .configServer(flag).replication(replication).cmdOptions(cmdOptions)
          .build();
      return mongodConfig;
    }

    private static IMongosConfig createIMongosConfig()
        throws UnknownHostException, IOException {
      IMongoCmdOptions cmdOptions = new MongoCmdOptionsBuilder().useNoJournal(false).verbose(false)
          .build();

      IMongosConfig mongosConfig = new MongosConfigBuilder()
          .version(Version.Main.PRODUCTION)
          .net(new Net(LOCALHOST, MONGOS_PORT, Network.localhostIsIPv6()))
          .configDB(LOCALHOST + ":" + CONFIG_SERVER_PORT)
          .cmdOptions(cmdOptions).build();
      return mongosConfig;
    }

    private static void cleanup() {
      if (mongosTestFactory != null) {
        mongosTestFactory.stop();
      }
    }

  }

  private static class SingleMode {

    private static MongodExecutable mongodExecutable;
    private static MongodProcess mongod;

    private static void setup() throws UnknownHostException, IOException {
      IMongoCmdOptions cmdOptions = new MongoCmdOptionsBuilder().verbose(false)
          .enableAuth(authEnabled).build();

      IMongodConfig mongodConfig = new MongodConfigBuilder()
          .version(Version.Main.PRODUCTION)
          .net(new Net(LOCALHOST, MONGOS_PORT, Network.localhostIsIPv6()))
          .cmdOptions(cmdOptions).build();

      IRuntimeConfig runtimeConfig = new RuntimeConfigBuilder().defaults(
          Command.MongoD).build();
      mongodExecutable = MongodStarter.getInstance(runtimeConfig).prepare(
          mongodConfig);
      mongod = mongodExecutable.start();
      mongoClient = new MongoClient(new ServerAddress(LOCALHOST, MONGOS_PORT));
      createDbAndCollections(EMPLOYEE_DB, EMPINFO_COLLECTION, "employee_id");
      createDbAndCollections(EMPLOYEE_DB, SCHEMA_CHANGE_COLLECTION, "field_2");
      createDbAndCollections(EMPLOYEE_DB, EMPTY_COLLECTION, "field_2");
      createDbAndCollections(DATATYPE_DB, DATATYPE_COLLECTION, "_id");
    }

    private static void cleanup() {
      if (mongod != null) {
        mongod.stop();
      }
      if (mongodExecutable != null) {
        mongodExecutable.stop();
      }
    }
  }

  @BeforeClass
  public static void initMongo() throws Exception {
    synchronized (MongoTestSuit.class) {
      if (initCount.get() == 0) {
        if (distMode) {
          logger.info("Executing tests in distributed mode");
          DistributedMode.setup();
        } else {
          logger.info("Executing tests in single mode");
          SingleMode.setup();
        }
        TestTableGenerator.importData(EMPLOYEE_DB, EMPINFO_COLLECTION, EMP_DATA);
        TestTableGenerator.importData(EMPLOYEE_DB, SCHEMA_CHANGE_COLLECTION, SCHEMA_CHANGE_DATA);
        TestTableGenerator.importData(DONUTS_DB, DONUTS_COLLECTION, DONUTS_DATA);
        TestTableGenerator.importData(DATATYPE_DB, DATATYPE_COLLECTION, DATATYPE_DATA);
      }
      initCount.incrementAndGet();
    }
  }

  private static void createDbAndCollections(String dbName,
      String collectionName, String indexFieldName) {
    MongoDatabase db = mongoClient.getDatabase(dbName);
    MongoCollection<Document> mongoCollection = db
        .getCollection(collectionName);
    if (mongoCollection == null) {
      db.createCollection(collectionName);
      mongoCollection = db.getCollection(collectionName);
    }
    IndexOptions indexOptions = new IndexOptions().unique(true)
        .background(false).name(indexFieldName);
    Bson keys = Indexes.ascending(indexFieldName);
    mongoCollection.createIndex(keys, indexOptions);
  }

  @AfterClass
  public static void tearDownCluster() throws Exception {
    synchronized (MongoTestSuit.class) {
      if (initCount.decrementAndGet() == 0) {
        try {
          if (mongoClient != null) {
            mongoClient.dropDatabase(EMPLOYEE_DB);
            mongoClient.dropDatabase(DATATYPE_DB);
            mongoClient.dropDatabase(DONUTS_DB);
          }
        }
        finally {
          if (mongoClient != null) {
            mongoClient.close();
          }
          if (distMode) {
            DistributedMode.cleanup();
          } else {
            SingleMode.cleanup();
          }
        }
      }
    }
  }
}