Merge pull request #517 from materialsproject/pymongo_updates

Changes to accommodate new pymongo

requirements.txt

@@ -1,4 +1,4 @@
-pymongo==3.11.2
+pymongo==4.0
 monty==2021.7.8
 mongomock==3.22.1
 pydash==4.9.3

setup.py

@@ -27,7 +27,7 @@
     include_package_data=True,
     install_requires=[
         "setuptools",
-        "pymongo>=3.6",
+        "pymongo>=4.0",
         "monty>=1.0.2",
         "mongomock>=3.10.0",
         "pydash>=4.1.0",

src/maggma/stores/mongolike.py

@@ -310,7 +310,7 @@ def count(self, criteria: Optional[Dict] = None) -> int:
         """
 
         criteria = criteria if criteria else {}
-        return self._collection.find(filter=criteria).count()
+        return self._collection.count_documents(filter=criteria)
 
     def query(  # type: ignore
         self,

tests/builders/test_simple_bib_drone.py

@@ -3,7 +3,7 @@
 
 import pytest
 
-from maggma.stores import MongoStore
+from maggma.stores import MemoryStore
 
 from .simple_bib_drone import SimpleBibDrone
 
@@ -16,9 +16,7 @@ def init_drone(test_dir):
     :return:
         initialized drone
     """
-    mongo_store = MongoStore(
-        database="drone_test", collection_name="drone_test", key="record_key"
-    )
+    mongo_store = MemoryStore(collection_name="drone_test", key="record_key")
     simple_path = test_dir / "simple_bib_example_data"
     assert simple_path.exists(), f"{simple_path} not found"
     simple_bib_drone = SimpleBibDrone(store=mongo_store, path=simple_path)

tests/stores/test_mongolike.py

@@ -168,7 +168,7 @@ def test_mongostore_from_db_file(mongostore, db_json):
 
 
 def test_mongostore_from_launchpad_file(lp_file):
-    ms = MongoStore.from_launchpad_file(lp_file, collection_name='tmp')
+    ms = MongoStore.from_launchpad_file(lp_file, collection_name="tmp")
     ms.connect()
     assert ms._collection.full_name == "maggma_tests.tmp"
 
@@ -288,30 +288,34 @@ def test_json_store_writeable(test_dir):
         jsonstore = JSONStore("d.json", file_writable=True)
         jsonstore.connect()
         assert jsonstore.count() == 2
-        jsonstore.update({'new': 'hello', 'task_id': 2})
+        jsonstore.update({"new": "hello", "task_id": 2})
         assert jsonstore.count() == 3
         jsonstore.close()
         jsonstore = JSONStore("d.json", file_writable=True)
         jsonstore.connect()
         assert jsonstore.count() == 3
-        jsonstore.remove_docs({'a': 5})
+        jsonstore.remove_docs({"a": 5})
         assert jsonstore.count() == 2
         jsonstore.close()
         jsonstore = JSONStore("d.json", file_writable=True)
         jsonstore.connect()
         assert jsonstore.count() == 2
         jsonstore.close()
-        with mock.patch("maggma.stores.JSONStore.update_json_file") as update_json_file_mock:
+        with mock.patch(
+            "maggma.stores.JSONStore.update_json_file"
+        ) as update_json_file_mock:
             jsonstore = JSONStore("d.json", file_writable=False)
             jsonstore.connect()
-            jsonstore.update({'new': 'hello', 'task_id': 5})
+            jsonstore.update({"new": "hello", "task_id": 5})
             assert jsonstore.count() == 3
             jsonstore.close()
             update_json_file_mock.assert_not_called()
-        with mock.patch("maggma.stores.JSONStore.update_json_file") as update_json_file_mock:
+        with mock.patch(
+            "maggma.stores.JSONStore.update_json_file"
+        ) as update_json_file_mock:
             jsonstore = JSONStore("d.json", file_writable=False)
             jsonstore.connect()
-            jsonstore.remove_docs({'task_id': 5})
+            jsonstore.remove_docs({"task_id": 5})
             assert jsonstore.count() == 2
             jsonstore.close()
             update_json_file_mock.assert_not_called()