class TestBoltztrap4DosBuilder(unittest.TestCase): def setUp(self): self.materials = JSONStore(boltztrap4dos_mat) self.materials.connect() self.bandstructure = JSONStore(boltztrap4dos_bs) self.bandstructure.connect() self.dos_ref = JSONStore(boltztrap4dos_dos) self.dos_ref.connect() self.dos = MemoryStore("dos") self.dos.connect() @unittest.skipIf("TRAVIS" in os.environ and os.environ["TRAVIS"] == "true", "Skipping this test on Travis CI.") def test_process_items(self): dosbuilder = Boltztrap4DosBuilder(self.materials, self.bandstructure, self.dos, avoid_projections=True) item = self.materials.query_one() bs_dict = self.bandstructure.query_one() item["bandstructure_uniform"] = bs_dict dos = dosbuilder.process_item(item) density = dos['densities']['1'][3900] self.assertAlmostEqual(density, 5.446126162946311, 5) def test_update_targets(self): dos = self.dos_ref.query_one() items = [dos] dosbuilder = Boltztrap4DosBuilder(self.materials, self.bandstructure, self.dos) dosbuilder.update_targets(items) self.assertListEqual(self.dos.distinct("task_id"), ['mp-663338'])
class TaskTaggerTest(unittest.TestCase): def setUp(self): # Set up test db, set up mpsft, etc. self.test_tasks = JSONStore(test_tasks) self.task_types = MemoryStore("task_types") self.test_tasks.connect() self.task_types.connect() def test_mp_defs(self): task_tagger = TaskTagger(tasks=self.test_tasks, task_types=self.task_types) for t in task_tagger.get_items(): processed = task_tagger.process_item(t) true_type = self.test_tasks.query_one( criteria={"task_id": t["task_id"]}, properties=["true_task_type"], )["true_task_type"] self.assertEqual(processed["task_type"], true_type)