Repository: Epistimio/orion

Issue:
Adjust `orion db setup` prompts based on the storage backend used.
The prompts are tailored for MongoDB specifically and misleading for pickleddb. The series of prompts should be adjusted based on which storage backend is specified at the first prompt.

Discussion:

Patch:
diff --git a/src/orion/core/cli/db/setup.py b/src/orion/core/cli/db/setup.py
index b129a814..7e98ab4d 100644
--- a/src/orion/core/cli/db/setup.py
+++ b/src/orion/core/cli/db/setup.py
@@ -14,6 +14,7 @@ import os
 import yaml
 
 import orion.core
+from orion.core.io.database import Database
 from orion.core.utils.terminal import ask_question
 
 log = logging.getLogger(__name__)
@@ -53,11 +54,23 @@ def main(*args):
         if cancel.strip().lower() == "n":
             return
 
-    _type = ask_question("Enter the database type: ", "mongodb")
-    name = ask_question("Enter the database name: ", "test")
-    host = ask_question("Enter the database host: ", "localhost")
-
-    config = {"database": {"type": _type, "name": name, "host": host}}
+    # Get database type.
+    _type = ask_question(
+        "Enter the database",
+        choice=Database.typenames,
+        default="mongodb",
+        ignore_case=True,
+    ).lower()
+    # Get database arguments.
+    db_class = Database.types[Database.typenames.index(_type)]
+    db_args = db_class.get_defaults()
+    arg_vals = {}
+    for arg_name, default_value in sorted(db_args.items()):
+        arg_vals[arg_name] = ask_question(
+            "Enter the database {}: ".format(arg_name), default_value
+        )
+
+    config = {"database": {"type": _type, **arg_vals}}
 
     print("Default configuration file will be saved at: ")
     print(default_file)
diff --git a/src/orion/core/io/database/__init__.py b/src/orion/core/io/database/__init__.py
index a5d87565..19bdc6fa 100644
--- a/src/orion/core/io/database/__init__.py
+++ b/src/orion/core/io/database/__init__.py
@@ -49,6 +49,10 @@ class AbstractDB(object, metaclass=AbstractSingletonType):
         **kwargs
     ):
         """Init method, see attributes of :class:`AbstractDB`."""
+        defaults = self.get_defaults()
+        host = defaults.get("host", None) if host is None or host == "" else host
+        name = defaults.get("name", None) if name is None or name == "" else name
+
         self.host = host
         self.name = name
         self.port = port
@@ -266,6 +270,20 @@ class AbstractDB(object, metaclass=AbstractSingletonType):
         """
         pass
 
+    @classmethod
+    @abstractmethod
+    def get_defaults(cls):
+        """Get database arguments needed to create a database instance.
+
+        Returns
+        -------
+        dict
+            A dictionary mapping an argument name to a default value.
+            If unexpected, default value can be None.
+
+        """
+        pass
+
 
 # pylint: disable=too-few-public-methods
 class ReadOnlyDB(object):
diff --git a/src/orion/core/io/database/ephemeraldb.py b/src/orion/core/io/database/ephemeraldb.py
index a3b07b98..f8d20bae 100644
--- a/src/orion/core/io/database/ephemeraldb.py
+++ b/src/orion/core/io/database/ephemeraldb.py
@@ -130,6 +130,16 @@ class EphemeralDB(AbstractDB):
 
         return dbcollection.delete_many(query=query)
 
+    @classmethod
+    def get_defaults(cls):
+        """Get database arguments needed to create a database instance.
+
+        .. seealso:: :meth:`orion.core.io.database.AbstractDB.get_defaults`
+                     for argument documentation.
+
+        """
+        return {}
+
 
 class EphemeralCollection(object):
     """Non permanent collection
diff --git a/src/orion/core/io/database/mongodb.py b/src/orion/core/io/database/mongodb.py
index 99749212..bf945278 100644
--- a/src/orion/core/io/database/mongodb.py
+++ b/src/orion/core/io/database/mongodb.py
@@ -330,3 +330,13 @@ class MongoDB(AbstractDB):
             self.options["authSource"] = settings["options"].get(
                 "authsource", self.name
             )
+
+    @classmethod
+    def get_defaults(cls):
+        """Get database arguments needed to create a database instance.
+
+        .. seealso:: :meth:`orion.core.io.database.AbstractDB.get_defaults`
+                     for argument documentation.
+
+        """
+        return {"name": "orion", "host": "localhost"}
diff --git a/src/orion/core/io/database/pickleddb.py b/src/orion/core/io/database/pickleddb.py
index f838dbf2..a67f9828 100644
--- a/src/orion/core/io/database/pickleddb.py
+++ b/src/orion/core/io/database/pickleddb.py
@@ -239,6 +239,16 @@ class PickledDB(AbstractDB):
         except Timeout as e:
             raise DatabaseTimeout(TIMEOUT_ERROR_MESSAGE.format(self.timeout)) from e
 
+    @classmethod
+    def get_defaults(cls):
+        """Get database arguments needed to create a database instance.
+
+        .. seealso:: :meth:`orion.core.io.database.AbstractDB.get_defaults`
+                     for argument documentation.
+
+        """
+        return {"host": DEFAULT_HOST}
+
 
 local_file_systems = ["ext2", "ext3", "ext4", "ntfs"]
 
diff --git a/src/orion/core/utils/terminal.py b/src/orion/core/utils/terminal.py
index db098854..0dbfbe00 100644
--- a/src/orion/core/utils/terminal.py
+++ b/src/orion/core/utils/terminal.py
@@ -6,7 +6,7 @@ Helper functions for terminal i/o
 """
 
 
-def ask_question(question, default=None):
+def ask_question(question, default=None, choice=None, ignore_case=False):
     """Ask a question to the user and receive an answer.
 
     Parameters
@@ -15,6 +15,11 @@ def ask_question(question, default=None):
         The question to be asked.
     default: str
         The default value to use if the user enters nothing.
+    choice: list
+        List of expected values to check user answer
+    ignore_case: bool
+        Used only if choice is provided. If True, ignore case when checking
+        user answer against given choice.
 
     Returns
     -------
@@ -22,13 +27,28 @@ def ask_question(question, default=None):
         The answer provided by the user.
 
     """
+    if choice is not None:
+        if ignore_case:
+            choice = [value.lower() for value in choice]
+        question = question + " (choice: {})".format(", ".join(choice))
+
     if default is not None:
         question = question + " (default: {}) ".format(default)
 
-    answer = input(question)
-
-    if answer.strip() == "":
-        return default
+    while True:
+        answer = input(question)
+        if answer.strip() == "":
+            answer = default
+            break
+        if choice is None:
+            break
+        if answer in choice or (ignore_case and answer.lower() in choice):
+            break
+        print(
+            "Unexpected value: {}. Must be one of: {}\n".format(
+                answer, ", ".join(choice)
+            )
+        )
 
     return answer
 


Test patch:
diff --git a/tests/functional/commands/test_setup_command.py b/tests/functional/commands/test_setup_command.py
index 7a4b105c..803a0480 100644
--- a/tests/functional/commands/test_setup_command.py
+++ b/tests/functional/commands/test_setup_command.py
@@ -8,6 +8,7 @@ import yaml
 
 import orion.core
 import orion.core.cli
+from orion.core.io.database import Database
 
 
 class _mock_input:
@@ -26,7 +27,7 @@ def test_creation_when_not_existing(monkeypatch, tmp_path):
     """Test if a configuration file is created when it does not exist."""
     config_path = str(tmp_path) + "/tmp_config.yaml"
     monkeypatch.setattr(orion.core, "DEF_CONFIG_FILES_PATHS", [config_path])
-    monkeypatch.setattr(builtins, "input", _mock_input(["type", "name", "host"]))
+    monkeypatch.setattr(builtins, "input", _mock_input(["mongodb", "host", "name"]))
 
     try:
         os.remove(config_path)
@@ -40,14 +41,16 @@ def test_creation_when_not_existing(monkeypatch, tmp_path):
     with open(config_path, "r") as output:
         content = yaml.safe_load(output)
 
-    assert content == {"database": {"type": "type", "name": "name", "host": "host"}}
+    assert content == {"database": {"type": "mongodb", "name": "name", "host": "host"}}
 
 
 def test_creation_when_exists(monkeypatch, tmp_path):
     """Test if the configuration file is overwritten when it exists."""
     config_path = str(tmp_path) + "/tmp_config.yaml"
     monkeypatch.setattr(orion.core, "DEF_CONFIG_FILES_PATHS", [config_path])
-    monkeypatch.setattr(builtins, "input", _mock_input(["y", "type", "name", "host"]))
+    monkeypatch.setattr(
+        builtins, "input", _mock_input(["y", "mongodb", "host", "name"])
+    )
 
     dump = {"database": {"type": "allo2", "name": "allo2", "host": "allo2"}}
 
@@ -81,6 +84,48 @@ def test_stop_creation_when_exists(monkeypatch, tmp_path):
     assert content == dump
 
 
+def test_invalid_database(monkeypatch, tmp_path, capsys):
+    """Test if command prompt loops when invalid database is typed."""
+    invalid_db_names = [
+        "invalid database",
+        "invalid database again",
+        "2383ejdd",
+        "another invalid database",
+    ]
+    config_path = str(tmp_path) + "/tmp_config.yaml"
+    monkeypatch.setattr(orion.core, "DEF_CONFIG_FILES_PATHS", [config_path])
+    monkeypatch.setattr(
+        builtins,
+        "input",
+        _mock_input(
+            [
+                *invalid_db_names,
+                "mongodb",
+                "the host",
+                "the name",
+            ]
+        ),
+    )
+
+    orion.core.cli.main(["db", "setup"])
+
+    with open(config_path, "r") as output:
+        content = yaml.safe_load(output)
+
+    assert content == {
+        "database": {"type": "mongodb", "name": "the name", "host": "the host"}
+    }
+
+    captured_output = capsys.readouterr().out
+    for invalid_db_name in invalid_db_names:
+        assert (
+            "Unexpected value: {}. Must be one of: {}\n".format(
+                invalid_db_name, ", ".join(Database.typenames)
+            )
+            in captured_output
+        )
+
+
 def test_defaults(monkeypatch, tmp_path):
     """Test if the default values are used when nothing user enters nothing."""
     config_path = str(tmp_path) + "/tmp_config.yaml"
@@ -93,5 +138,35 @@ def test_defaults(monkeypatch, tmp_path):
         content = yaml.safe_load(output)
 
     assert content == {
-        "database": {"type": "mongodb", "name": "test", "host": "localhost"}
+        "database": {"type": "mongodb", "name": "orion", "host": "localhost"}
     }
+
+
+def test_ephemeraldb(monkeypatch, tmp_path):
+    """Test if config content is written for an ephemeraldb."""
+    config_path = str(tmp_path) + "/tmp_config.yaml"
+    monkeypatch.setattr(orion.core, "DEF_CONFIG_FILES_PATHS", [config_path])
+    monkeypatch.setattr(builtins, "input", _mock_input(["ephemeraldb"]))
+
+    orion.core.cli.main(["db", "setup"])
+
+    with open(config_path, "r") as output:
+        content = yaml.safe_load(output)
+
+    assert content == {"database": {"type": "ephemeraldb"}}
+
+
+def test_pickleddb(monkeypatch, tmp_path):
+    """Test if config content is written for an pickleddb."""
+    host = "my_pickles.db"
+
+    config_path = str(tmp_path) + "/tmp_config.yaml"
+    monkeypatch.setattr(orion.core, "DEF_CONFIG_FILES_PATHS", [config_path])
+    monkeypatch.setattr(builtins, "input", _mock_input(["pickleddb", host]))
+
+    orion.core.cli.main(["db", "setup"])
+
+    with open(config_path, "r") as output:
+        content = yaml.safe_load(output)
+
+    assert content == {"database": {"type": "pickleddb", "host": host}}