From 95435efa49c6d20c40db0ff135ce30c1e3b9285d Mon Sep 17 00:00:00 2001
From: Daniel Hollas <daniel.hollas@bristol.ac.uk>
Date: Thu, 21 Nov 2024 17:36:04 +0100
Subject: [PATCH] Use enum

---
 tests/conftest.py | 50 +++++++++++++++++++++++++++++++----------------
 1 file changed, 33 insertions(+), 17 deletions(-)

diff --git a/tests/conftest.py b/tests/conftest.py
index 6df03794f..841693d04 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -22,6 +22,7 @@
 import types
 import typing as t
 import warnings
+from enum import Enum
 from pathlib import Path
 
 import click
@@ -37,6 +38,13 @@
 pytest_plugins = ['aiida.tools.pytest_fixtures', 'sphinx.testing.fixtures']
 
 
+class DbBackend(Enum):
+    """Options for the '--db-backend' CLI argument"""
+
+    SQLITE = 'sqlite'
+    PSQL = 'psql'
+
+
 def pytest_collection_modifyitems(items, config):
     """Automatically generate markers for certain tests.
 
@@ -47,12 +55,13 @@ def pytest_collection_modifyitems(items, config):
     filepath_django = Path(__file__).parent / 'storage' / 'psql_dos' / 'migrations' / 'django_branch'
     filepath_sqla = Path(__file__).parent / 'storage' / 'psql_dos' / 'migrations' / 'sqlalchemy_branch'
 
-    if config.option.db_backend == 'sqlite':
-        # TODO: Validate that user did not do '--db-backend sqlite -m requires_psql'
-        if config.option.markexpr == '':
-            config.option.markexpr = 'not requires_psql'
-        else:
+    # If the user requested the SQlite backend, automatically skip incompatible tests
+    if config.option.db_backend is DbBackend.SQLITE:
+        if config.option.markexpr != '':
+            # Don't overwrite markers that the user already provided via '-m ' cmdline argument
             config.option.markexpr += ' and not requires_psql'
+        else:
+            config.option.markexpr = 'not requires_psql'
 
     for item in items:
         filepath_item = Path(item.fspath)
@@ -75,24 +84,27 @@ def pytest_collection_modifyitems(items, config):
             item.add_marker('presto')
 
 
-def db_backend_checker(value):
-    """Validation function for the custom '--db-backend' pytest CLI option"""
-    # TODO: Use enum
-    db_backends = ('sqlite', 'psql')
-    if value not in db_backends:
-        msg = f"Invalid database backend option '{value}', must be one of: {db_backends}"
+def db_backend_type(string):
+    """Conversion function for the custom '--db-backend' pytest CLI option
+
+    :param string: String provided by the user via CLI
+    :returns: DbBackend enum corresponding to user input string
+    """
+    try:
+        return DbBackend(string)
+    except ValueError:
+        msg = f"Invalid --db-backend option '{string}'\nMust be one of: {tuple(db.value for db in DbBackend)}"
         raise pytest.UsageError(msg)
-    return value
 
 
 def pytest_addoption(parser):
     parser.addoption(
         '--db-backend',
         action='store',
-        default='sqlite',
+        default=DbBackend.SQLITE,
         required=False,
-        help="Database backend to be used for tests ('sqlite' or 'psql')",
-        type=db_backend_checker,
+        help=f'Database backend to be used for tests {tuple(db.value for db in DbBackend)}',
+        type=db_backend_type,
     )
 
 
@@ -105,18 +117,22 @@ def aiida_profile(pytestconfig, aiida_config, aiida_profile_factory, config_psql
     be run against the main storage backend, which is ``core.sqlite_dos``.
     """
     marker_opts = pytestconfig.getoption('-m')
+    db_backend = pytestconfig.getoption('--db-backend')
 
     # By default we use RabbitMQ broker and psql_dos storage
     broker = 'core.rabbitmq'
     if 'not requires_rmq' in marker_opts or 'presto' in marker_opts:
         broker = None
 
-    if 'not requires_psql' in marker_opts or 'presto' in marker_opts:
+    if db_backend is DbBackend.SQLITE:
         storage = 'core.sqlite_dos'
         config = config_sqlite_dos()
-    else:
+    elif db_backend is DbBackend.PSQL:
         storage = 'core.psql_dos'
         config = config_psql_dos()
+    else:
+        # This should be unreachable
+        raise ValueError(f'Invalid DB backend {db_backend}')
 
     with aiida_profile_factory(
         aiida_config, storage_backend=storage, storage_config=config, broker_backend=broker