[WIP] Remove MongoDB dependency from core tests

src/orion/core/io/database/pickleddb.py

@@ -214,7 +214,7 @@ def _dump_database(self, database):
             log.error(
                 "Document in (collection: %s) is not pickable\ndoc: %s",
                 collection,
-                doc.to_dict(),
+                doc.to_dict() if hasattr(doc, "to_dict") else str(doc),
             )
 
             key, value = find_unpickable_field(doc)

src/orion/storage/legacy.py

@@ -355,7 +355,7 @@ def count_broken_trials(self, experiment):
     def update_heartbeat(self, trial):
         """Update trial's heartbeat"""
         return self.update_trial(
-            trial, heartbeat=datetime.datetime.utcnow(), status="reserved"
+            trial, heartbeat=datetime.datetime.utcnow(), where={"status": "reserved"}
         )
 
     def fetch_trials_by_status(self, experiment, status):

tests/conftest.py

@@ -19,7 +19,7 @@
 from orion.core.io.database.pickleddb import PickledDB
 from orion.core.utils.singleton import update_singletons
 from orion.core.worker.trial import Trial
-from orion.storage.base import Storage
+from orion.storage.base import Storage, get_storage, setup_storage
 from orion.storage.legacy import Legacy
 from orion.testing import OrionState
 
@@ -315,24 +315,6 @@ def update_metadata(metadata):
     monkeypatch.setattr(resolve_config, "update_metadata", update_metadata)
 
 
-@pytest.fixture()
-def create_db_instance(null_db_instances, clean_db):
-    """Create and save a singleton database instance."""
-    try:
-        database = {
-            "type": "MongoDB",
-            "name": "orion_test",
-            "username": "user",
-            "password": "pass",
-        }
-        db = Storage(of_type="legacy", database=database)
-    except ValueError:
-        db = Storage()
-
-    db = db._db
-    return db
-
-
 @pytest.fixture()
 def script_path():
     """Return a script path for mock"""
@@ -372,6 +354,12 @@ def setup_pickleddb_database():
     del os.environ["ORION_DB_ADDRESS"]
 
 
+@pytest.fixture(scope="function")
+def storage(setup_pickleddb_database):
+    setup_storage()
+    yield get_storage()
+
+
 @pytest.fixture()
 def with_user_userxyz(monkeypatch):
     """Make ``getpass.getuser()`` return ``'userxyz'``."""

tests/functional/branching/orion_config.yaml

@@ -7,8 +7,3 @@ algorithms: random
 
 producer:
     strategy: NoParallelStrategy
-
-database:
-  type: 'mongodb'
-  name: 'orion_test'
-  host: 'mongodb://user:pass@localhost'

tests/functional/branching/test_branching.py

@@ -18,7 +18,7 @@ def execute(command, assert_code=0):
 
 
 @pytest.fixture
-def init_full_x(create_db_instance, monkeypatch):
+def init_full_x(setup_pickleddb_database, monkeypatch):
     """Init original experiment"""
     monkeypatch.chdir(os.path.dirname(os.path.abspath(__file__)))
     name = "full_x"
@@ -984,7 +984,9 @@ def test_auto_resolution_with_fidelity(init_full_x_full_y, monkeypatch):
     ]
 
 
-def test_init_w_version_from_parent_w_children(create_db_instance, monkeypatch, capsys):
+def test_init_w_version_from_parent_w_children(
+    setup_pickleddb_database, monkeypatch, capsys
+):
     """Test that init of experiment from version with children fails."""
     monkeypatch.chdir(os.path.dirname(os.path.abspath(__file__)))
     execute(
@@ -1007,7 +1009,7 @@ def test_init_w_version_from_parent_w_children(create_db_instance, monkeypatch,
     assert "Experiment name" in captured.err
 
 
-def test_init_w_version_from_exp_wout_child(create_db_instance, monkeypatch):
+def test_init_w_version_from_exp_wout_child(setup_pickleddb_database, monkeypatch):
     """Test that init of experiment from version without child works."""
     monkeypatch.chdir(os.path.dirname(os.path.abspath(__file__)))
     execute(
@@ -1026,7 +1028,7 @@ def test_init_w_version_from_exp_wout_child(create_db_instance, monkeypatch):
     assert len(list(exp))
 
 
-def test_init_w_version_gt_max(create_db_instance, monkeypatch):
+def test_init_w_version_gt_max(setup_pickleddb_database, monkeypatch):
     """Test that init of experiment from version higher than max works."""
     monkeypatch.chdir(os.path.dirname(os.path.abspath(__file__)))
     execute(
@@ -1045,7 +1047,7 @@ def test_init_w_version_gt_max(create_db_instance, monkeypatch):
     assert len(list(exp))
 
 
-def test_init_check_increment_w_children(create_db_instance, monkeypatch):
+def test_init_check_increment_w_children(setup_pickleddb_database, monkeypatch):
     """Test that incrementing version works with not same-named children."""
     monkeypatch.chdir(os.path.dirname(os.path.abspath(__file__)))
     execute(
@@ -1064,7 +1066,7 @@ def test_init_check_increment_w_children(create_db_instance, monkeypatch):
     assert len(list(exp))
 
 
-def test_branch_from_selected_version(create_db_instance, monkeypatch):
+def test_branch_from_selected_version(setup_pickleddb_database, monkeypatch):
     """Test that branching from a version passed with `--version` works."""
     monkeypatch.chdir(os.path.dirname(os.path.abspath(__file__)))
     execute(

tests/functional/client/test_cli_client.py

@@ -9,9 +9,7 @@
 from orion.core.worker.consumer import Consumer
 
 
-@pytest.mark.usefixtures("clean_db")
-@pytest.mark.usefixtures("null_db_instances")
-def test_interrupt(database, monkeypatch, capsys):
+def test_interrupt(storage, monkeypatch, capsys):
     """Test interruption from within user script."""
     monkeypatch.chdir(os.path.dirname(os.path.abspath(__file__)))
 
@@ -37,17 +35,15 @@ def test_interrupt(database, monkeypatch, capsys):
     assert captured.out == "Orion is interrupted.\n"
     assert captured.err == ""
 
-    exp = list(database.experiments.find({"name": "voila_voici"}))
+    exp = list(storage.fetch_experiments({"name": "voila_voici"}))
     exp = exp[0]
     exp_id = exp["_id"]
-    trials = list(database.trials.find({"experiment": exp_id}))
+    trials = list(storage.fetch_trials(uid=exp_id))
     assert len(trials) == 1
-    assert trials[0]["status"] == "interrupted"
+    assert trials[0].status == "interrupted"
 
 
-@pytest.mark.usefixtures("clean_db")
-@pytest.mark.usefixtures("null_db_instances")
-def test_interrupt_diff_code(database, monkeypatch, capsys):
+def test_interrupt_diff_code(storage, monkeypatch, capsys):
     """Test interruption from within user script with custom int code"""
     monkeypatch.chdir(os.path.dirname(os.path.abspath(__file__)))
 
@@ -80,12 +76,12 @@ def empty_env(self, trial, results_file=None):
 
         assert error_code == 0
 
-        exp = list(database.experiments.find({"name": "voila_voici"}))
+        exp = list(storage.fetch_experiments({"name": "voila_voici"}))
         exp = exp[0]
         exp_id = exp["_id"]
-        trials = list(database.trials.find({"experiment": exp_id}))
+        trials = list(storage.fetch_trials(uid=exp_id))
         assert len(trials) == 2
-        assert trials[0]["status"] == "broken"
+        assert trials[0].status == "broken"
 
     # This time we use true `get_execution_environment which pass properly int code to child.
     error_code = orion.core.cli.main(
@@ -108,12 +104,12 @@ def empty_env(self, trial, results_file=None):
     assert "Orion is interrupted.\n" in captured.out
     assert captured.err == ""
 
-    exp = list(database.experiments.find({"name": "voila_voici"}))
+    exp = list(storage.fetch_experiments({"name": "voila_voici"}))
     exp = exp[0]
     exp_id = exp["_id"]
-    trials = list(database.trials.find({"experiment": exp_id}))
+    trials = list(storage.fetch_trials(uid=exp_id))
     assert len(trials) == 3
-    assert trials[-1]["status"] == "interrupted"
+    assert trials[-1].status == "interrupted"
 
 
 # TODO:
@@ -124,9 +120,7 @@ def empty_env(self, trial, results_file=None):
 
 
 @pytest.mark.parametrize("fct", ["report_bad_trial", "report_objective"])
-@pytest.mark.usefixtures("clean_db")
-@pytest.mark.usefixtures("null_db_instances")
-def test_report_no_name(database, monkeypatch, fct):
+def test_report_no_name(storage, monkeypatch, fct):
     """Test report helper functions with default names"""
     monkeypatch.chdir(os.path.dirname(os.path.abspath(__file__)))
 
@@ -148,21 +142,19 @@ def test_report_no_name(database, monkeypatch, fct):
         + user_args
     )
 
-    exp = list(database.experiments.find({"name": "voila_voici"}))
+    exp = list(storage.fetch_experiments({"name": "voila_voici"}))
     exp = exp[0]
     exp_id = exp["_id"]
-    trials = list(database.trials.find({"experiment": exp_id}))
+    trials = list(storage.fetch_trials(uid=exp_id))
     assert len(trials) == 2
-    assert trials[0]["status"] == "completed"
-    assert trials[0]["results"][0]["name"] == "objective"
-    assert trials[0]["results"][0]["type"] == "objective"
-    assert trials[0]["results"][0]["value"] == 1.0
+    assert trials[0].status == "completed"
+    assert trials[0].results[0].name == "objective"
+    assert trials[0].results[0].type == "objective"
+    assert trials[0].results[0].value == 1.0
 
 
 @pytest.mark.parametrize("fct", ["report_bad_trial", "report_objective"])
-@pytest.mark.usefixtures("clean_db")
-@pytest.mark.usefixtures("null_db_instances")
-def test_report_with_name(database, monkeypatch, fct):
+def test_report_with_name(storage, monkeypatch, fct):
     """Test report helper functions with custom names"""
     monkeypatch.chdir(os.path.dirname(os.path.abspath(__file__)))
 
@@ -186,21 +178,19 @@ def test_report_with_name(database, monkeypatch, fct):
         + user_args
     )
 
-    exp = list(database.experiments.find({"name": "voila_voici"}))
+    exp = list(storage.fetch_experiments({"name": "voila_voici"}))
     exp = exp[0]
     exp_id = exp["_id"]
-    trials = list(database.trials.find({"experiment": exp_id}))
+    trials = list(storage.fetch_trials(uid=exp_id))
     assert len(trials) == 2
-    assert trials[0]["status"] == "completed"
-    assert trials[0]["results"][0]["name"] == "metric"
-    assert trials[0]["results"][0]["type"] == "objective"
-    assert trials[0]["results"][0]["value"] == 1.0
+    assert trials[0].status == "completed"
+    assert trials[0].results[0].name == "metric"
+    assert trials[0].results[0].type == "objective"
+    assert trials[0].results[0].value == 1.0
 
 
 @pytest.mark.parametrize("fct", ["report_bad_trial", "report_objective"])
-@pytest.mark.usefixtures("clean_db")
-@pytest.mark.usefixtures("null_db_instances")
-def test_report_with_bad_objective(database, monkeypatch, fct):
+def test_report_with_bad_objective(storage, monkeypatch, fct):
     """Test report helper functions with bad objective types"""
     monkeypatch.chdir(os.path.dirname(os.path.abspath(__file__)))
 
@@ -226,9 +216,7 @@ def test_report_with_bad_objective(database, monkeypatch, fct):
     assert "must contain a type `objective` with type float/int" in str(exc.value)
 
 
-@pytest.mark.usefixtures("clean_db")
-@pytest.mark.usefixtures("null_db_instances")
-def test_report_with_bad_trial_no_objective(database, monkeypatch):
+def test_report_with_bad_trial_no_objective(storage, monkeypatch):
     """Test bad trial report helper function with default objective."""
     monkeypatch.chdir(os.path.dirname(os.path.abspath(__file__)))
 
@@ -248,20 +236,18 @@ def test_report_with_bad_trial_no_objective(database, monkeypatch):
         + user_args
     )
 
-    exp = list(database.experiments.find({"name": "voila_voici"}))
+    exp = list(storage.fetch_experiments({"name": "voila_voici"}))
     exp = exp[0]
     exp_id = exp["_id"]
-    trials = list(database.trials.find({"experiment": exp_id}))
+    trials = list(storage.fetch_trials(uid=exp_id))
     assert len(trials) == 2
-    assert trials[0]["status"] == "completed"
-    assert trials[0]["results"][0]["name"] == "objective"
-    assert trials[0]["results"][0]["type"] == "objective"
-    assert trials[0]["results"][0]["value"] == 1e10
+    assert trials[0].status == "completed"
+    assert trials[0].results[0].name == "objective"
+    assert trials[0].results[0].type == "objective"
+    assert trials[0].results[0].value == 1e10
 
 
-@pytest.mark.usefixtures("clean_db")
-@pytest.mark.usefixtures("null_db_instances")
-def test_report_with_bad_trial_with_data(database, monkeypatch):
+def test_report_with_bad_trial_with_data(storage, monkeypatch):
     """Test bad trial report helper function with additional data."""
     monkeypatch.chdir(os.path.dirname(os.path.abspath(__file__)))
 
@@ -283,24 +269,22 @@ def test_report_with_bad_trial_with_data(database, monkeypatch):
         + user_args
     )
 
-    exp = list(database.experiments.find({"name": "voila_voici"}))
+    exp = list(storage.fetch_experiments({"name": "voila_voici"}))
     exp = exp[0]
     exp_id = exp["_id"]
-    trials = list(database.trials.find({"experiment": exp_id}))
+    trials = list(storage.fetch_trials(uid=exp_id))
     assert len(trials) == 2
-    assert trials[0]["status"] == "completed"
-    assert trials[0]["results"][0]["name"] == "objective"
-    assert trials[0]["results"][0]["type"] == "objective"
-    assert trials[0]["results"][0]["value"] == 1e10
+    assert trials[0].status == "completed"
+    assert trials[0].results[0].name == "objective"
+    assert trials[0].results[0].type == "objective"
+    assert trials[0].results[0].value == 1e10
 
-    assert trials[0]["results"][1]["name"] == "another"
-    assert trials[0]["results"][1]["type"] == "constraint"
-    assert trials[0]["results"][1]["value"] == 1.0
+    assert trials[0].results[1].name == "another"
+    assert trials[0].results[1].type == "constraint"
+    assert trials[0].results[1].value == 1.0
 
 
-@pytest.mark.usefixtures("clean_db")
-@pytest.mark.usefixtures("null_db_instances")
-def test_no_report(database, monkeypatch, capsys):
+def test_no_report(storage, monkeypatch, capsys):
     """Test script call without any results reported."""
     monkeypatch.chdir(os.path.dirname(os.path.abspath(__file__)))