Reuse functionality between AiidaTestCase and pytest fixtures

aiida/backends/testbase.py

@@ -12,7 +12,7 @@
 import unittest
 import traceback
 
-from aiida.common.exceptions import ConfigurationError, TestsNotAllowedError, InternalError
+from aiida.common.exceptions import TestsNotAllowedError
 from aiida.manage import configuration
 from aiida.manage.manager import get_manager, reset_manager
 from aiida import orm
@@ -33,49 +33,17 @@ class AiidaTestCase(unittest.TestCase):
 
     Internally it loads the AiidaTestImplementation subclass according to the current backend."""
     _computer = None  # type: aiida.orm.Computer
-    _user = None  # type: aiida.orm.User
     _class_was_setup = False
-    __backend_instance = None
     backend = None  # type: aiida.orm.implementation.Backend
 
-    @classmethod
-    def get_backend_class(cls):
-        """Get backend class."""
-        from aiida.backends.testimplbase import AiidaTestImplementation
-        from aiida.backends import BACKEND_SQLA, BACKEND_DJANGO
-        from aiida.manage.configuration import PROFILE
-
-        # Freeze the __impl_class after the first run
-        if not hasattr(cls, '__impl_class'):
-            if PROFILE.database_backend == BACKEND_SQLA:
-                from aiida.backends.sqlalchemy.testbase import SqlAlchemyTests
-                cls.__impl_class = SqlAlchemyTests
-            elif PROFILE.database_backend == BACKEND_DJANGO:
-                from aiida.backends.djsite.db.testbase import DjangoTests
-                cls.__impl_class = DjangoTests
-            else:
-                raise ConfigurationError('Unknown backend type')
-
-            # Check that it is of the right class
-            if not issubclass(cls.__impl_class, AiidaTestImplementation):
-                raise InternalError(
-                    'The AiiDA test implementation is not of type '
-                    '{}, that is not a subclass of AiidaTestImplementation'.format(cls.__impl_class.__name__)
-                )
-
-        return cls.__impl_class
-
     @classmethod
     def setUpClass(cls):
         """Set up test class."""
-        # Note: this will raise an exception, that will be seen as a test
-        # failure. To be safe, you should do the same check also in the tearDownClass
-        # to avoid that it is run
+        # Note: this will raise an exception, that will be seen as a test failure.
         check_if_tests_can_run()
 
         # Force the loading of the backend which will load the required database environment
         cls.backend = get_manager().get_backend()
-        cls.__backend_instance = cls.get_backend_class()()
         cls._class_was_setup = True
 
         cls.refurbish_db()
@@ -84,33 +52,13 @@ def setUpClass(cls):
     def tearDownClass(cls):
         """Tear down test class.
 
-        Note: Also cleans file repository.
+        Note: Cleans database and also the file repository.
         """
-        # Double check for double security to avoid to run the tearDown
-        # if this is not a test profile
-
-        check_if_tests_can_run()
-        if orm.autogroup.CURRENT_AUTOGROUP is not None:
-            orm.autogroup.CURRENT_AUTOGROUP.clear_group_cache()
         cls.clean_db()
         cls.clean_repository()
 
-    def tearDown(self):
-        reset_manager()
-
     ### Database/repository-related methods
 
-    @classmethod
-    def insert_data(cls):
-        """
-        This method setups the database (by creating a default user) and
-        inserts default data into the database (which is for the moment a
-        default computer).
-        """
-        orm.User.objects.reset()  # clear Aiida's cache of the default user
-        # populate user cache of test clases
-        cls.user  # pylint: disable=pointless-statement
-
     @classmethod
     def clean_db(cls):
         """Clean up database and reset caches.
@@ -119,17 +67,17 @@ def clean_db(cls):
         """
         from aiida.common.exceptions import InvalidOperation
 
-        # Note: this will raise an exception, that will be seen as a test
-        # failure. To be safe, you should do the same check also in the tearDownClass
-        # to avoid that it is run
+        # Note: this will raise an exception, that will be seen as a test failure.
+        # Just another safety check to prevent deleting production databases
         check_if_tests_can_run()
 
         if not cls._class_was_setup:
             raise InvalidOperation('You cannot call clean_db before running the setUpClass')
 
-        cls.__backend_instance.clean_db()
+        cls.backend._clean_db()  # pylint: disable=protected-access
         cls._computer = None
-        cls._user = None
+
+        orm.User.objects.reset()  # clear Aiida's cache of the default user
 
         if orm.autogroup.CURRENT_AUTOGROUP is not None:
             orm.autogroup.CURRENT_AUTOGROUP.clear_group_cache()
@@ -138,12 +86,14 @@ def clean_db(cls):
 
     @classmethod
     def refurbish_db(cls):
-        """Clean up database and repopulate with initial data.
+        """Clean up database and create default user.
 
-        Combines clean_db and insert_data.
+        Combines clean_db and database initialization.
         """
         cls.clean_db()
-        cls.insert_data()
+        created, user = orm.User.objects.get_or_create_default()  # create default user
+        if created:
+            user.store()
 
     @classmethod
     def clean_repository(cls):
@@ -190,10 +140,12 @@ def computer(cls):  # pylint: disable=no-self-argument
         return cls._computer
 
     @classproperty
-    def user(cls):  # pylint: disable=no-self-argument
-        if cls._user is None:
-            cls._user = get_default_user()
-        return cls._user
+    def user(cls):  # pylint: disable=no-self-argument,no-self-use
+        """Return default user.
+
+        Since the default user is already cached at the orm level, we just return it.
+        """
+        return orm.User.objects.get_default()
 
     @classproperty
     def user_email(cls):  # pylint: disable=no-self-argument
@@ -224,27 +176,3 @@ def tearDownClass(cls, *args, **kwargs):
         """Close the PGTest postgres test cluster."""
         super().tearDownClass(*args, **kwargs)
         cls.pg_test.close()
-
-
-def get_default_user(**kwargs):
-    """Creates and stores the default user in the database.
-
-    Default user email is taken from current profile.
-    No-op if user already exists.
-    The same is done in `verdi setup`.
-
-    :param kwargs: Additional information to use for new user, i.e. 'first_name', 'last_name' or 'institution'.
-    :returns: the :py:class:`~aiida.orm.User`
-    """
-    from aiida.manage.configuration import get_config
-    email = get_config().current_profile.default_user
-
-    if kwargs.pop('email', None):
-        raise ValueError('Do not specify the user email (must coincide with default user email of profile).')
-
-    # Create the AiiDA user if it does not yet exist
-    created, user = orm.User.objects.get_or_create(email=email, **kwargs)
-    if created:
-        user.store()
-
-    return user

aiida/manage/tests/__init__.py

@@ -137,46 +137,22 @@ def __init__(self, profile_name):
             raise TestManagerError('Unable to load test profile \'{}\'.'.format(profile_name))
         check_if_tests_can_run()
 
-        self._select_db_test_case(backend=self._profile.database_backend)
-
-    def _select_db_test_case(self, backend):
-        """
-        Selects tests case for the correct database backend.
-        """
-        if backend == BACKEND_DJANGO:
-            from aiida.backends.djsite.db.testbase import DjangoTests
-            self._test_case = DjangoTests()
-        elif backend == BACKEND_SQLA:
-            from aiida.backends.sqlalchemy.testbase import SqlAlchemyTests
-            from aiida.backends.sqlalchemy import get_scoped_session
-
-            self._test_case = SqlAlchemyTests()
-            self._test_case.test_session = get_scoped_session()
-
     def reset_db(self):
-        self._test_case.clean_db()  # will drop all users
+        self._profile.database_backend._clean_db()  # pylint: disable=protected-access
         manager.reset_manager()
         self.init_db()
 
-    def init_db(self):
+    def init_db(self):  # pylint: disable=no-self-use
         """Initialise the database state for running of tests.
 
         Adds default user if necessary.
         """
         from aiida.orm import User
-        from aiida.cmdline.commands.cmd_user import set_default_user
 
         if not User.objects.get_default():
-            user_dict = get_user_dict(_DEFAULT_PROFILE_INFO)
-            try:
-                user = User(**user_dict)
-                user.store()
-            except exceptions.IntegrityError:
-                # The user already exists, no problem
-                user = User.objects.get(**user_dict)
-
-            set_default_user(self._profile, user)
             User.objects.reset()  # necessary to pick up new default user
+            _created, _user = User.objects.get_or_create_default(**get_user_dict(_DEFAULT_PROFILE_INFO))
+            _user.store()
 
     def has_profile_open(self):
         return self._profile is not None
@@ -231,8 +207,6 @@ class TemporaryProfileManager(ProfileManager):
 
     """
 
-    _test_case = None
-
     def __init__(self, backend=BACKEND_DJANGO, pgtest=None):  # pylint: disable=super-init-not-called
         """Construct a TemporaryProfileManager
 
@@ -340,7 +314,6 @@ def create_profile(self):
         backend = manager.get_manager()._load_backend(schema_check=False)
         backend.migrate()
 
-        self._select_db_test_case(backend=self._profile.database_backend)
         self.init_db()
 
     def repo_ok(self):

aiida/orm/implementation/django/backend.py

@@ -142,3 +142,16 @@ def get_connection():
         # For now we just return the global but if we ever support multiple Django backends
         # being loaded this should be specific to this backend
         return connection
+
+    def _clean_db(self):
+        from aiida.backends.djsite.db import models as dbmodels
+        # I first need to delete the links, because in principle I could not delete input nodes, only outputs.
+        # For simplicity, since I am deleting everything, I delete the links first
+        dbmodels.DbLink.objects.all().delete()
+
+        # Then I delete the nodes, otherwise I cannot delete computers and users
+        dbmodels.DbLog.objects.all().delete()
+        dbmodels.DbNode.objects.all().delete()  # pylint: disable=no-member
+        dbmodels.DbUser.objects.all().delete()  # pylint: disable=no-member
+        dbmodels.DbComputer.objects.all().delete()
+        dbmodels.DbGroup.objects.all().delete()

aiida/orm/implementation/sql/backends.py

@@ -74,3 +74,10 @@ def execute_prepared_statement(self, sql, parameters):
                 results.append(row)
 
         return results
+
+    @abc.abstractmethod
+    def _clean_db(self):
+        """Cleans entire database (used in tests).
+
+        .. warning:: this will lead to catastrophic data loss!
+        """

aiida/orm/implementation/sqlalchemy/backend.py

@@ -157,3 +157,26 @@ def get_connection():
         """
         from aiida.backends import sqlalchemy as sa
         return sa.ENGINE.raw_connection()
+
+    def _clean_db(self):
+        from sqlalchemy.sql import table
+        # pylint: disable=invalid-name
+        DbGroupNodes = table('db_dbgroup_dbnodes')
+        DbGroup = table('db_dbgroup')
+        DbLink = table('db_dblink')
+        DbNode = table('db_dbnode')
+        DbLog = table('db_dblog')
+        DbAuthInfo = table('db_dbauthinfo')
+        DbUser = table('db_dbuser')
+        DbComputer = table('db_dbcomputer')
+
+        with self.transaction() as session:
+            session.execute(DbGroupNodes.delete())
+            session.execute(DbGroup.delete())
+            session.execute(DbLog.delete())
+            session.execute(DbLink.delete())
+            session.execute(DbNode.delete())
+            session.execute(DbAuthInfo.delete())
+            session.execute(DbComputer.delete())
+            session.execute(DbUser.delete())
+            session.commit()