style cleanup

materialsproject · Mar 9, 2018 · f8d7d8b · f8d7d8b
1 parent a0df960
commit f8d7d8b
Show file tree

Hide file tree

Showing 5 changed files with 34 additions and 61 deletions.
diff --git a/emmet/vasp/builders/tests/test_elastic.py b/emmet/vasp/builders/tests/test_elastic.py
@@ -11,8 +11,8 @@
 __email__ = "[email protected]"
 
 module_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)))
-test_tasks = os.path.join(module_dir, "..","..","..", "..", "test_files",
-                          "vasp", "elastic_tasks.json")
+test_tasks = os.path.join(module_dir, "..", "..", "..", "..", "test_files", "vasp", "elastic_tasks.json")
+
 
 class ElasticBuilderTest(unittest.TestCase):
     @classmethod
@@ -27,24 +27,26 @@ def setUpClass(cls):
         # Generate test materials collection
         cls.test_materials = MongoStore("test_emmet", "materials")
         cls.test_materials.connect()
-        opt_docs = cls.test_tasks.query(["output.structure", "formula_pretty"],
-                                        {"task_label": "structure optimization"})
-        mat_docs = [{"material_id": "mp-{}".format(n),
-                     "structure": opt_doc['output']['structure'],
-                     "pretty_formula": opt_doc['formula_pretty']}
-                    for n, opt_doc in enumerate(opt_docs)]
+        opt_docs = cls.test_tasks.query(["output.structure", "formula_pretty"], {
+            "task_label": "structure optimization"
+        })
+        mat_docs = [{
+            "material_id": "mp-{}".format(n),
+            "structure": opt_doc['output']['structure'],
+            "pretty_formula": opt_doc['formula_pretty']
+        } for n, opt_doc in enumerate(opt_docs)]
         cls.test_materials.update(mat_docs, key='material_id', update_lu=False)
 
     def test_builder(self):
-        ec_builder = ElasticBuilder(self.test_tasks, self.test_elasticity,
-                                    self.test_materials, incremental=False)
+        ec_builder = ElasticBuilder(self.test_tasks, self.test_elasticity, self.test_materials, incremental=False)
         ec_builder.connect()
         for t in ec_builder.get_items():
             processed = ec_builder.process_item(t)
             if processed:
                 pass
             else:
-                import nose; nose.tools.set_trace()
+                import nose
+                nose.tools.set_trace()
         runner = Runner([ec_builder])
         runner.run()
 

diff --git a/emmet/vasp/builders/tests/test_electronic_structure.py b/emmet/vasp/builders/tests/test_electronic_structure.py
@@ -10,7 +10,6 @@
 
 
 class TestElectronicStructureImageBuilder(unittest.TestCase):
-
     def setUp(self):
 
         coords = list()
@@ -28,8 +27,7 @@ def setUp(self):
     def test_get_bandstructure(self):
 
         self.builder.bfs = MagicMock()
-        mat = {"bandstructure": {"bs_oid": "234234",
-                                 "bs_compression": "zlib"}}
+        mat = {"bandstructure": {"bs_oid": "234234", "bs_compression": "zlib"}}
 
         with patch("emmet.vasp.builders.electronic_structure_img.json") as json_patch:
             with patch("emmet.vasp.builders.electronic_structure_img.zlib") as zlib_patch:
@@ -42,8 +40,7 @@ def test_get_bandstructure(self):
     def test_get_dos(self):
 
         self.builder.dfs = MagicMock()
-        mat = {"bandstructure": {"dos_oid": "234234",
-                                 "dos_compression": "zlib"}}
+        mat = {"bandstructure": {"dos_oid": "234234", "dos_compression": "zlib"}}
 
         with patch("emmet.vasp.builders.electronic_structure_img.json") as json_patch:
             with patch("emmet.vasp.builders.electronic_structure_img.zlib") as zlib_patch:

diff --git a/emmet/vasp/builders/tests/test_ml_structure.py b/emmet/vasp/builders/tests/test_ml_structure.py
@@ -2,7 +2,6 @@
 from pymatgen.core.structure import Structure
 from pymatgen.core.lattice import Lattice
 
-
 from maggma.stores import MongoStore
 from emmet.vasp.builders.ml_structures import MLStructuresBuilder
 
@@ -11,7 +10,6 @@
 
 
 class TestML(unittest.TestCase):
-
     def setUp(self):
         tasks = MongoStore("emmet_test", "tasks")
         ml_strucs = MongoStore("emmet_test", "ml_strucs")
@@ -46,10 +44,7 @@ def test_process_item(self):
         coords[1][0] = 0.0
         structure2 = Structure(lattice, ["Si", "Si"], coords)
 
-        dummy_task["calcs_reversed"][0]["input"] = {"incar": {
-            "IBRION": 2,
-            "ISIF": 3,
-            "NSW": 99}}
+        dummy_task["calcs_reversed"][0]["input"] = {"incar": {"IBRION": 2, "ISIF": 3, "NSW": 99}}
 
         ionic_steps = dummy_task["calcs_reversed"][0]["output"]["ionic_steps"]
         ionic_steps.append({"structure": structure1.as_dict()})

diff --git a/emmet/vasp/builders/tests/test_mpworks.py b/emmet/vasp/builders/tests/test_mpworks.py
@@ -9,8 +9,7 @@
 __email__ = "[email protected]"
 
 module_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)))
-test_tasks = os.path.join(module_dir, "..", "..", "..", "..",
-                          "test_files", "mpworks_tasks.json")
+test_tasks = os.path.join(module_dir, "..", "..", "..", "..", "test_files", "mpworks_tasks.json")
 
 
 class MPWorksCompatibilityBuilderTest(unittest.TestCase):
@@ -22,9 +21,7 @@ def setUp(self):
         self.elasticity.connect()
 
     def test_builder(self):
-        mpw_builder = MPWorksCompatibilityBuilder(self.test_tasks,
-                                                  self.elasticity,
-                                                  incremental=False)
+        mpw_builder = MPWorksCompatibilityBuilder(self.test_tasks, self.elasticity, incremental=False)
         items = mpw_builder.get_items()
         processed = [mpw_builder.process_item(item) for item in items]
         mpw_builder.update_targets(processed)
@@ -37,12 +34,10 @@ def test_set_mongolike(self):
         self.assertEqual(test_dict["new"]["value"], 5)
 
     def test_convert_mpworks_to_atomate(self):
-        doc = self.test_tasks.collection.find_one(
-            {"task_type": {"$regex": "deformed"}})
+        doc = self.test_tasks.collection.find_one({"task_type": {"$regex": "deformed"}})
         new_doc = convert_mpworks_to_atomate(doc)
         self.assertTrue('hubbards' in new_doc['input'])
-        doc = self.test_tasks.collection.find_one(
-            {"task_type": {"$regex": "(2x)"}})
+        doc = self.test_tasks.collection.find_one({"task_type": {"$regex": "(2x)"}})
         new_doc = convert_mpworks_to_atomate(doc)
         self.assertTrue('hubbards' in new_doc['input'])
 

diff --git a/emmet/vasp/builders/tests/test_structure_similarity.py b/emmet/vasp/builders/tests/test_structure_similarity.py
@@ -2,67 +2,51 @@
 import os
 
 from emmet.vasp.builders.structure_similarity import *
-from maggma.stores import MongoStore
+from maggma.stores import MemoryStore, JSONStore
 
 from monty.serialization import loadfn
 
 __author__ = "Nils E. R. Zimmermann"
 __email__ = "[email protected]"
 
 module_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)))
-test_site_fp_stats = os.path.join(module_dir, "..","..","..", "..", "test_files",
-                                  "site_fingerprint_stats.json")
+test_site_fp_stats = os.path.join(module_dir, "..", "..", "..", "..", "test_files", "site_fingerprint_stats.json")
+
 
 class StructureSimilarityBuilderTest(unittest.TestCase):
     def setUp(self):
         # Set up test db, etc.
-        self.test_site_descriptors = MongoStore("test_emmet", 
-                                                "struct_sim_site_fp",
-                                                key="task_id")
+        self.test_site_descriptors = JSONStore(test_site_fp_stats)
         self.test_site_descriptors.connect()
-        self.site_fp_docs = loadfn(test_site_fp_stats, cls=None)
-        self.test_site_descriptors.update(self.site_fp_docs)
-        self.test_structure_similarity = MongoStore("test_emmet",
-                                                    "struct_sim_struct_sim",
-                                                    "task_ids")
+        self.test_structure_similarity = MemoryStore("struct_sim_struct_sim")
 
     def test_builder(self):
-        sim_builder = StructureSimilarityBuilder(
-                self.test_site_descriptors,
-                self.test_structure_similarity)
+        sim_builder = StructureSimilarityBuilder(self.test_site_descriptors, self.test_structure_similarity)
         sim_builder.connect()
         for t in sim_builder.get_items():
             processed = sim_builder.process_item(t)
             if processed:
                 pass
             else:
-                import nose; nose.tools.set_trace()
+                import nose
+                nose.tools.set_trace()
 
     def test_get_all_site_descriptors(self):
-        sim_builder = StructureSimilarityBuilder(
-                self.test_site_descriptors,
-                self.test_structure_similarity)
+        sim_builder = StructureSimilarityBuilder(self.test_site_descriptors, self.test_structure_similarity)
         for i in range(3):
-            d = sim_builder.get_similarities(
-                    self.site_fp_docs[i],
-                    self.site_fp_docs[i])
+            d = sim_builder.get_similarities(self.site_fp_docs[i], self.site_fp_docs[i])
             self.assertAlmostEqual(d['cos'], 1)
             self.assertAlmostEqual(d['dist'], 0)
-        d = sim_builder.get_similarities(
-                self.site_fp_docs[0],
-                self.site_fp_docs[1])
+        d = sim_builder.get_similarities(self.site_fp_docs[0], self.site_fp_docs[1])
         self.assertAlmostEqual(d['cos'], 0.0013649)
         self.assertAlmostEqual(d['dist'], 2.6866749)
-        d = sim_builder.get_similarities(
-                self.site_fp_docs[0],
-                self.site_fp_docs[2])
+        d = sim_builder.get_similarities(self.site_fp_docs[0], self.site_fp_docs[2])
         self.assertAlmostEqual(d['cos'], 0.0013069)
         self.assertAlmostEqual(d['dist'], 2.6293889)
-        d = sim_builder.get_similarities(
-                self.site_fp_docs[1],
-                self.site_fp_docs[2])
+        d = sim_builder.get_similarities(self.site_fp_docs[1], self.site_fp_docs[2])
         self.assertAlmostEqual(d['cos'], 0.0012729)
         self.assertAlmostEqual(d['dist'], 2.7235044)
 
+
 if __name__ == "__main__":
     unittest.main()