def test_assimilate(self):
        simulate = True if VaspToDbTaskDroneTest.conn is None else False
        drone = VaspToDbTaskDrone(database="creator_unittest",
                                  simulate_mode=simulate)
        queen = BorgQueen(drone)
        queen.serial_assimilate(os.path.join(test_dir, 'db_test'))
        data = queen.get_data()
        self.assertEqual(len(data), 5)
        if VaspToDbTaskDroneTest.conn:
            db = VaspToDbTaskDroneTest.conn["creator_unittest"]
            data = db.tasks.find()
            self.assertEqual(data.count(), 5)
            warnings.warn("Actual db insertion mode.")

        for d in data:
            dir_name = d['dir_name']
            if dir_name.endswith("killed_mp_aflow"):
                self.assertEqual(d['state'], "killed")
                self.assertFalse(d['is_hubbard'])
                self.assertEqual(d['pretty_formula'], "SiO2")
            elif dir_name.endswith("stopped_mp_aflow"):
                self.assertEqual(d['state'], "stopped")
                self.assertEqual(d['pretty_formula'], "ThFe5P3")
            elif dir_name.endswith("success_mp_aflow"):
                self.assertEqual(d['state'], "successful")
                self.assertEqual(d['pretty_formula'], "TbZn(BO2)5")
                self.assertAlmostEqual(d['output']['final_energy'],
                                       -526.66747274, 4)
            elif dir_name.endswith("Li2O_aflow"):
                self.assertEqual(d['state'], "successful")
                self.assertEqual(d['pretty_formula'], "Li2O")
                self.assertAlmostEqual(d['output']['final_energy'],
                                       -14.31446494, 6)
                self.assertEqual(len(d["calculations"]), 2)
            elif dir_name.endswith("Li2O"):
                self.assertEqual(d['state'], "successful")
                self.assertEqual(d['pretty_formula'], "Li2O")
                self.assertAlmostEqual(d['output']['final_energy'],
                                       -14.31337758, 6)
                self.assertEqual(len(d["calculations"]), 1)
                self.assertEqual(len(d["custodian"]), 1)
                self.assertEqual(len(d["custodian"][0]["corrections"]), 1)

        if VaspToDbTaskDroneTest.conn:
            warnings.warn("Testing query engine mode.")
            qe = QueryEngine(database="creator_unittest")
            self.assertEqual(qe.query().count(), 5)
            #Test mappings by query engine.
            for r in qe.query(criteria={"pretty_formula": "Li2O"},
                              properties=["dir_name", "energy",
                                          "calculations"]):
                if r["dir_name"].endswith("Li2O_aflow"):
                    self.assertAlmostEqual(r['energy'], -14.31446494, 4)
                    self.assertEqual(len(r["calculations"]), 2)
                elif r["dir_name"].endswith("Li2O"):
                    self.assertAlmostEqual(r['energy'],
                                           -14.31337758, 4)
                    self.assertEqual(len(r["calculations"]), 1)

            # Test query one.
            d = qe.query_one(criteria={"pretty_formula": "TbZn(BO2)5"},
                             properties=["energy"])
            self.assertAlmostEqual(d['energy'], -526.66747274, 4)

            d = qe.get_entries_in_system(["Li", "O"])
            self.assertEqual(len(d), 2)
            self.assertIsInstance(d[0], ComputedEntry)

            s = qe.get_structure_from_id(d[0].entry_id)
            self.assertIsInstance(s, Structure)
            self.assertEqual(s.formula, "Li2 O1")
    def test_assimilate(self):
        """Borg assimilation code.
        This takes too long for a unit test!
        """
        simulate = True if VaspToDbTaskDroneTest.conn is None else False
        drone = VaspToDbTaskDrone(database="creator_unittest",
                                  simulate_mode=simulate,
                                  parse_dos=True,
                                  compress_dos=1)
        queen = BorgQueen(drone)
        queen.serial_assimilate(os.path.join(test_dir, 'db_test'))
        data = queen.get_data()
        self.assertEqual(len(data), 6)
        if VaspToDbTaskDroneTest.conn:
            db = VaspToDbTaskDroneTest.conn["creator_unittest"]
            data = db.tasks.find()
            self.assertEqual(data.count(), 6)
            warnings.warn("Actual db insertion mode.")

        for d in data:
            dir_name = d['dir_name']
            if dir_name.endswith("killed_mp_aflow"):
                self.assertEqual(d['state'], "killed")
                self.assertFalse(d['is_hubbard'])
                self.assertEqual(d['pretty_formula'], "SiO2")
            elif dir_name.endswith("stopped_mp_aflow"):
                self.assertEqual(d['state'], "stopped")
                self.assertEqual(d['pretty_formula'], "ThFe5P3")
            elif dir_name.endswith("success_mp_aflow"):
                self.assertEqual(d['state'], "successful")
                self.assertEqual(d['pretty_formula'], "TbZn(BO2)5")
                self.assertAlmostEqual(d['output']['final_energy'],
                                       -526.66747274, 4)
            elif dir_name.endswith("Li2O_aflow"):
                self.assertEqual(d['state'], "successful")
                self.assertEqual(d['pretty_formula'], "Li2O")
                self.assertAlmostEqual(d['output']['final_energy'],
                                       -14.31446494, 6)
                self.assertEqual(len(d["calculations"]), 2)
                self.assertEqual(d['input']['is_lasph'], False)
                self.assertEqual(d['input']['xc_override'], None)
            elif dir_name.endswith("Li2O"):
                self.assertEqual(d['state'], "successful")
                self.assertEqual(d['pretty_formula'], "Li2O")
                self.assertAlmostEqual(d['output']['final_energy'],
                                       -14.31337758, 6)
                self.assertEqual(len(d["calculations"]), 1)
                self.assertEqual(len(d["custodian"]), 1)
                self.assertEqual(len(d["custodian"][0]["corrections"]), 1)
            elif dir_name.endswith("Li2O_aflow_lasph"):
                self.assertEqual(d['state'], "successful")
                self.assertEqual(d['pretty_formula'], "Li2O")
                self.assertAlmostEqual(d['output']['final_energy'], -13.998171,
                                       6)
                self.assertEqual(len(d["calculations"]), 2)
                self.assertEqual(d['input']['is_lasph'], True)
                self.assertEqual(d['input']['xc_override'], "PS")

        if VaspToDbTaskDroneTest.conn:
            warnings.warn("Testing query engine mode.")
            qe = QueryEngine(database="creator_unittest")
            self.assertEqual(qe.query().count(), 6)
            #Test mappings by query engine.
            for r in qe.query(
                    criteria={"pretty_formula": "Li2O"},
                    properties=["dir_name", "energy", "calculations",
                                "input"]):
                if r["dir_name"].endswith("Li2O_aflow"):
                    self.assertAlmostEqual(r['energy'], -14.31446494, 4)
                    self.assertEqual(len(r["calculations"]), 2)
                    self.assertEqual(r["input"]["is_lasph"], False)
                    self.assertEqual(r['input']['xc_override'], None)
                elif r["dir_name"].endswith("Li2O"):
                    self.assertAlmostEqual(r['energy'], -14.31337758, 4)
                    self.assertEqual(len(r["calculations"]), 1)
                    self.assertEqual(r["input"]["is_lasph"], False)
                    self.assertEqual(r['input']['xc_override'], None)

            #Test lasph
            e = qe.get_entries({"dir_name": {"$regex": "lasph"}})
            self.assertEqual(len(e), 1)
            self.assertEqual(e[0].parameters["is_lasph"], True)
            self.assertEqual(e[0].parameters["xc_override"], "PS")

            # Test query one.
            d = qe.query_one(criteria={"pretty_formula": "TbZn(BO2)5"},
                             properties=["energy"])
            self.assertAlmostEqual(d['energy'], -526.66747274, 4)

            d = qe.get_entries_in_system(["Li", "O"])
            self.assertEqual(len(d), 3)
            self.assertIsInstance(d[0], ComputedEntry)

            s = qe.get_structure_from_id(d[0].entry_id)
            self.assertIsInstance(s, Structure)
            self.assertEqual(s.formula, "Li2 O1")

            self.assertIsInstance(qe.get_dos_from_id(d[0].entry_id),
                                  CompleteDos)
    def test_assimilate(self):
        """Borg assimilation code.
        This takes too long for a unit test!
        """
        simulate = True if VaspToDbTaskDroneTest.conn is None else False
        drone = VaspToDbTaskDrone(database="creator_unittest",
                                  simulate_mode=simulate,
                                  parse_dos=True, compress_dos=1)
        queen = BorgQueen(drone)
        queen.serial_assimilate(os.path.join(test_dir, 'db_test'))
        data = queen.get_data()
        self.assertEqual(len(data), 6)
        if VaspToDbTaskDroneTest.conn:
            db = VaspToDbTaskDroneTest.conn["creator_unittest"]
            data = db.tasks.find()
            self.assertEqual(data.count(), 6)
            warnings.warn("Actual db insertion mode.")

        for d in data:
            dir_name = d['dir_name']
            if dir_name.endswith("killed_mp_aflow"):
                self.assertEqual(d['state'], "killed")
                self.assertFalse(d['is_hubbard'])
                self.assertEqual(d['pretty_formula'], "SiO2")
            elif dir_name.endswith("stopped_mp_aflow"):
                self.assertEqual(d['state'], "stopped")
                self.assertEqual(d['pretty_formula'], "ThFe5P3")
            elif dir_name.endswith("success_mp_aflow"):
                self.assertEqual(d['state'], "successful")
                self.assertEqual(d['pretty_formula'], "TbZn(BO2)5")
                self.assertAlmostEqual(d['output']['final_energy'],
                                       -526.66747274, 4)
            elif dir_name.endswith("Li2O_aflow"):
                self.assertEqual(d['state'], "successful")
                self.assertEqual(d['pretty_formula'], "Li2O")
                self.assertAlmostEqual(d['output']['final_energy'],
                                       -14.31446494, 6)
                self.assertEqual(len(d["calculations"]), 2)
                self.assertEqual(d['input']['is_lasph'], False)
                self.assertEqual(d['input']['xc_override'], None)
                self.assertEqual(d["oxide_type"], "oxide")
            elif dir_name.endswith("Li2O"):
                self.assertEqual(d['state'], "successful")
                self.assertEqual(d['pretty_formula'], "Li2O")
                self.assertAlmostEqual(d['output']['final_energy'],
                                       -14.31337758, 6)
                self.assertEqual(len(d["calculations"]), 1)
                self.assertEqual(len(d["custodian"]), 1)
                self.assertEqual(len(d["custodian"][0]["corrections"]), 1)
            elif dir_name.endswith("Li2O_aflow_lasph"):
                self.assertEqual(d['state'], "successful")
                self.assertEqual(d['pretty_formula'], "Li2O")
                self.assertAlmostEqual(d['output']['final_energy'],
                                       -13.998171, 6)
                self.assertEqual(len(d["calculations"]), 2)
                self.assertEqual(d['input']['is_lasph'], True)
                self.assertEqual(d['input']['xc_override'], "PS")

        if VaspToDbTaskDroneTest.conn:
            warnings.warn("Testing query engine mode.")
            qe = QueryEngine(database="creator_unittest")
            self.assertEqual(qe.query().count(), 6)
            #Test mappings by query engine.
            for r in qe.query(criteria={"pretty_formula": "Li2O"},
                              properties=["dir_name", "energy",
                                          "calculations", "input"]):
                if r["dir_name"].endswith("Li2O_aflow"):
                    self.assertAlmostEqual(r['energy'], -14.31446494, 4)
                    self.assertEqual(len(r["calculations"]), 2)
                    self.assertEqual(r["input"]["is_lasph"], False)
                    self.assertEqual(r['input']['xc_override'], None)
                    self.assertEqual(d["oxide_type"], "oxide")
                elif r["dir_name"].endswith("Li2O"):
                    self.assertAlmostEqual(r['energy'],
                                           -14.31337758, 4)
                    self.assertEqual(len(r["calculations"]), 1)
                    self.assertEqual(r["input"]["is_lasph"], False)
                    self.assertEqual(r['input']['xc_override'], None)

            # Test lasph
            e = qe.get_entries({"dir_name":{"$regex":"lasph"}})
            self.assertEqual(len(e), 1)
            self.assertEqual(e[0].parameters["is_lasph"], True)
            self.assertEqual(e[0].parameters["xc_override"], "PS")

            # Test query one.
            d = qe.query_one(criteria={"pretty_formula": "TbZn(BO2)5"},
                             properties=["energy"])
            self.assertAlmostEqual(d['energy'], -526.66747274, 4)

            d = qe.get_entries_in_system(["Li", "O"])
            self.assertEqual(len(d), 3)
            self.assertIsInstance(d[0], ComputedEntry)
            self.assertEqual(d[0].data["oxide_type"], "oxide")

            s = qe.get_structure_from_id(d[0].entry_id)
            self.assertIsInstance(s, Structure)
            self.assertEqual(s.formula, "Li2 O1")

            self.assertIsInstance(qe.get_dos_from_id(d[0].entry_id), CompleteDos)