Optionally use mongomock instead of pymongo/mongodb

docs_rst/config_tutorial.rst

@@ -76,6 +76,7 @@ A few basic parameters that can be tweaked are:
 * ``WEBSERVER_HOST: 127.0.0.1`` - the default host on which to run the web server
 * ``WEBSERVER_PORT: 5000`` - the default port on which to run the web server
 * ``QUEUE_JOBNAME_MAXLEN: 20`` - the max length of the job name to send to the queuing system (some queuing systems limit the size of job names)
+* ``MONGOMOCK_SERVERSTORE_FILE`` - path to a non-empty JSON file, if set then mongomock will be used instead of MongoDB; this file should be initialized with '{}'
 
 Parameters that you probably shouldn't change
 ---------------------------------------------

fireworks/core/launchpad.py

@@ -15,11 +15,12 @@
 from bson import ObjectId
 from monty.os.path import zpath
 from monty.serialization import loadfn
-from pymongo import ASCENDING, DESCENDING, MongoClient
+from pymongo import ASCENDING, DESCENDING
 from pymongo.errors import DocumentTooLarge
 from tqdm import tqdm
 
 from fireworks.core.firework import Firework, FWAction, Launch, Tracker, Workflow
+from fireworks.fw_config import MongoClient
 from fireworks.fw_config import (
     GRIDFS_FALLBACK_COLLECTION,
     LAUNCHPAD_LOC,
@@ -411,7 +412,7 @@ def bulk_add_wfs(self, wfs):
 
         """
         # Make all fireworks workflows
-        wfs = [Workflow.from_firework(wf) if isinstance(wf, Firework) else wf for wf in wfs]
+        wfs = [Workflow.from_Firework(wf) if isinstance(wf, Firework) else wf for wf in wfs]
 
         # Initialize new firework counter, starting from the next fw id
         total_num_fws = sum(len(wf) for wf in wfs)

fireworks/core/tests/test_launchpad.py

@@ -15,7 +15,6 @@
 
 import pytest
 from monty.os import cd
-from pymongo import MongoClient
 from pymongo import __version__ as PYMONGO_VERSION
 from pymongo.errors import OperationFailure
 
@@ -43,7 +42,7 @@ class AuthenticationTest(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
         try:
-            client = MongoClient()
+            client = fireworks.fw_config.MongoClient()
             client.not_the_admin_db.command("createUser", "myuser", pwd="mypassword", roles=["dbOwner"])
         except Exception:
             raise unittest.SkipTest("MongoDB is not running in localhost:27017! Skipping tests.")

fireworks/fw_config.py

@@ -5,6 +5,9 @@
 
 from monty.design_patterns import singleton
 from monty.serialization import dumpfn, loadfn
+import pymongo
+import mongomock
+import mongomock.gridfs
 
 __author__ = "Anubhav Jain"
 __copyright__ = "Copyright 2012, The Materials Project"
@@ -102,6 +105,12 @@
 # a dynamically generated document exceeds the 16MB limit. Functionality disabled if None.
 GRIDFS_FALLBACK_COLLECTION = "fw_gridfs"
 
+# path to a database file to use with mongomock, do not use mongomock if None
+MONGOMOCK_SERVERSTORE_FILE = None
+
+# default mongoclient class
+MongoClient = pymongo.MongoClient
+
 
 def override_user_settings() -> None:
     module_dir = os.path.dirname(os.path.abspath(__file__))
@@ -153,6 +162,15 @@ def override_user_settings() -> None:
             if len(m_paths) > 0:
                 globals()[k] = m_paths[0]
 
+    if 'MONGOMOCK_SERVERSTORE_FILE' in os.environ:
+        globals()['MONGOMOCK_SERVERSTORE_FILE'] = os.environ['MONGOMOCK_SERVERSTORE_FILE']
+    if globals()['MONGOMOCK_SERVERSTORE_FILE']:
+        if not os.environ.get('MONGOMOCK_SERVERSTORE_FILE'):
+            os.environ['MONGOMOCK_SERVERSTORE_FILE'] = globals()['MONGOMOCK_SERVERSTORE_FILE']
+        globals()['MongoClient'] = getattr(mongomock, 'MongoClient')
+        if globals()['GRIDFS_FALLBACK_COLLECTION']:
+            mongomock.gridfs.enable_gridfs_integration()
+
 
 override_user_settings()
 

fireworks/user_objects/firetasks/filepad_tasks.py

@@ -1,7 +1,11 @@
 import os
-
+import json
+from glob import glob
+from pymongo import DESCENDING
+from ruamel.yaml import YAML
 from fireworks.core.firework import FiretaskBase
 from fireworks.utilities.filepad import FilePad
+from fireworks.utilities.dict_mods import arrow_to_dot
 
 __author__ = "Kiran Mathew, Johannes Hoermann"
 __email__ = "[email protected], [email protected]"
@@ -28,7 +32,6 @@ class AddFilesTask(FiretaskBase):
     optional_params = ["identifiers", "directory", "filepad_file", "compress", "metadata"]
 
     def run_task(self, fw_spec):
-        from glob import glob
 
         directory = os.path.abspath(self.get("directory", "."))
 
@@ -143,19 +146,12 @@ class GetFilesByQueryTask(FiretaskBase):
     ]
 
     def run_task(self, fw_spec):
-        import json
-
-        import pymongo
-        from ruamel.yaml import YAML
-
-        from fireworks.utilities.dict_mods import arrow_to_dot
-
         fpad = get_fpad(self.get("filepad_file", None))
         dest_dir = self.get("dest_dir", os.path.abspath("."))
         new_file_names = self.get("new_file_names", [])
         query = self.get("query", {})
         sort_key = self.get("sort_key", None)
-        sort_direction = self.get("sort_direction", pymongo.DESCENDING)
+        sort_direction = self.get("sort_direction", DESCENDING)
         limit = self.get("limit", None)
         fizzle_empty_result = self.get("fizzle_empty_result", True)
         fizzle_degenerate_file_name = self.get("fizzle_degenerate_file_name", True)

fireworks/utilities/filepad.py

@@ -7,11 +7,12 @@
 import zlib
 
 import gridfs
-import pymongo
+from pymongo import DESCENDING
 from monty.json import MSONable
 from monty.serialization import loadfn
-from pymongo import MongoClient
+from bson.objectid import ObjectId
 
+from fireworks.fw_config import MongoClient
 from fireworks.fw_config import LAUNCHPAD_LOC, MONGO_SOCKET_TIMEOUT_MS
 from fireworks.utilities.fw_utilities import get_fw_logger
 
@@ -176,7 +177,7 @@ def get_file_by_id(self, gfs_id):
         doc = self.filepad.find_one({"gfs_id": gfs_id})
         return self._get_file_contents(doc)
 
-    def get_file_by_query(self, query, sort_key=None, sort_direction=pymongo.DESCENDING):
+    def get_file_by_query(self, query, sort_key=None, sort_direction=DESCENDING):
         """
 
         Args:
@@ -289,8 +290,6 @@ def _get_file_contents(self, doc):
         Returns:
             (str, dict): the file content as a string, document dictionary
         """
-        from bson.objectid import ObjectId
-
         if doc:
             gfs_id = doc["gfs_id"]
             file_contents = self.gridfs.get(ObjectId(gfs_id)).read()

setup.py

@@ -39,6 +39,7 @@
             "tqdm>=4.8.4",
             "importlib-metadata>=4.8.2; python_version<'3.8'",
             "typing-extensions; python_version<'3.8'",
+            "mongomock @ git+https://github.com/ikondov/mongomock.git@persistence",
         ],
         extras_require={
             "rtransfer": ["paramiko>=2.4.2"],