Source code for aiida.manage.tests.main

# -*- coding: utf-8 -*-
###########################################################################
# Copyright (c), The AiiDA team. All rights reserved.                     #
# This file is part of the AiiDA code.                                    #
#                                                                         #
# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core #
# For further information on the license, see the LICENSE.txt file        #
# For further information please visit http://www.aiida.net               #
###########################################################################
"""
Testing infrastructure for easy testing of AiiDA plugins.

"""
from contextlib import contextmanager
import os
import shutil
import tempfile

from aiida.common.warnings import warn_deprecation
from aiida.manage import configuration, get_manager
from aiida.manage.configuration import settings
from aiida.manage.external.postgres import Postgres

__all__ = (
    'get_test_profile_name',
    'get_test_backend_name',
    'get_user_dict',
    'test_manager',
    'TestManager',
    'TestManagerError',
    'ProfileManager',
    'TemporaryProfileManager',
)

_DEFAULT_PROFILE_INFO = {
    'name': 'test_profile',
    'email': 'tests@aiida.mail',
    'first_name': 'AiiDA',
    'last_name': 'Plugintest',
    'institution': 'aiidateam',
    'storage_backend': 'core.psql_dos',
    'database_engine': 'postgresql_psycopg2',
    'database_username': 'aiida',
    'database_password': 'aiida_pw',
    'database_name': 'aiida_db',
    'repo_dir': 'test_repo',
    'config_dir': '.aiida',
    'root_path': '',
    'broker_protocol': 'amqp',
    'broker_username': 'guest',
    'broker_password': 'guest',
    'broker_host': '127.0.0.1',
    'broker_port': 5672,
    'broker_virtual_host': '',
    'test_profile': True,
}


[docs]class TestManagerError(Exception): """Raised by TestManager in situations that may lead to inconsistent behaviour."""
[docs] def __init__(self, msg): super().__init__() self.msg = msg
[docs] def __str__(self): return repr(self.msg)
[docs]class TestManager: """ Test manager for plugin tests. Uses either ProfileManager for wrapping an existing profile or TemporaryProfileManager for setting up a complete temporary AiiDA environment. For usage with pytest, see :py:class:`~aiida.manage.tests.pytest_fixtures`. """
[docs] def __init__(self): self._manager = None
@property def manager(self) -> 'ProfileManager': assert self._manager is not None return self._manager
[docs] def use_temporary_profile(self, backend=None, pgtest=None): """Set up Test manager to use temporary AiiDA profile. Uses :py:class:`aiida.manage.tests.main.TemporaryProfileManager` internally. :param backend: Backend to use. :param pgtest: a dictionary of arguments to be passed to PGTest() for starting the postgresql cluster, e.g. {'pg_ctl': '/somepath/pg_ctl'}. Should usually not be necessary. """ if configuration.get_profile() is not None: raise TestManagerError('An AiiDA profile must not be loaded before setting up a test profile.') if self._manager is not None: raise TestManagerError('Profile manager already loaded.') mngr = TemporaryProfileManager(backend=backend, pgtest=pgtest) mngr.create_profile() self._manager = mngr # don't assign before profile has actually been created!
[docs] def use_profile(self, profile_name): """Set up Test manager to use existing profile. Uses :py:class:`aiida.manage.tests.main.ProfileManager` internally. :param profile_name: Name of existing test profile to use. """ if configuration.get_profile() is not None: raise TestManagerError('an AiiDA profile must not be loaded before setting up a test profile.') if self._manager is not None: raise TestManagerError('Profile manager already loaded.') self._manager = ProfileManager(profile_name=profile_name)
[docs] def has_profile_open(self): return self._manager and self._manager.has_profile_open()
[docs] def reset_db(self): warn_deprecation('reset_db() is deprecated, use clear_profile() instead', version=3) return self._manager.clear_profile()
[docs] def clear_profile(self): """Reset the global profile, clearing all its data and closing any open resources.""" return self._manager.clear_profile()
[docs] def destroy_all(self): if self._manager: self._manager.destroy_all() self._manager = None
[docs]class ProfileManager: """ Wraps existing AiiDA profile. """
[docs] def __init__(self, profile_name): """ Use an existing profile. :param profile_name: Name of the profile to be loaded """ from aiida import load_profile self._profile = None try: self._profile = load_profile(profile_name) except Exception: raise TestManagerError(f'Unable to load test profile `{profile_name}`.') if self._profile is None: raise TestManagerError(f'Unable to load test profile `{profile_name}`.') if not self._profile.is_test_profile: raise TestManagerError(f'Profile `{profile_name}` is not a valid test profile.')
[docs] @staticmethod def clear_profile(): """Reset the global profile, clearing all its data and closing any open resources.""" manager = get_manager() manager.get_profile_storage()._clear(recreate_user=True) # pylint: disable=protected-access manager.get_profile_storage() # reload the storage connection manager.reset_communicator() manager.reset_runner()
[docs] def has_profile_open(self): return self._profile is not None
[docs] def destroy_all(self): # pylint: disable=no-self-use manager = get_manager() manager.reset_profile()
[docs]class TemporaryProfileManager(ProfileManager): """ Manage the life cycle of a completely separated and temporary AiiDA environment. * No profile / database setup required * Tests run via the TemporaryProfileManager never pollute the user's working environment Filesystem: * temporary ``.aiida`` configuration folder * temporary repository folder Database: * temporary database cluster (via the ``pgtest`` package) * with ``aiida`` database user * with ``aiida_db`` database AiiDA: * configured to use the temporary configuration * sets up a temporary profile for tests All of this happens automatically when using the corresponding tests classes & tests runners (unittest) or fixtures (pytest). Example:: tests = TemporaryProfileManager(backend=backend) tests.create_aiida_db() # set up only the database tests.create_profile() # set up a profile (creates the db too if necessary) # ready for tests # run tests 1 tests.clear_profile() # database ready for independent tests 2 # run tests 2 tests.destroy_all() # everything cleaned up """
[docs] def __init__(self, backend='core.psql_dos', pgtest=None): # pylint: disable=super-init-not-called """Construct a TemporaryProfileManager :param backend: a database backend :param pgtest: a dictionary of arguments to be passed to PGTest() for starting the postgresql cluster, e.g. {'pg_ctl': '/somepath/pg_ctl'}. Should usually not be necessary. """ self.dbinfo = {} self.profile_info = _DEFAULT_PROFILE_INFO self.profile_info['storage_backend'] = backend self._pgtest = pgtest or {} self.pg_cluster = None self.postgres = None self._profile = None self._has_test_db = False self._backup = { 'config': configuration.CONFIG, 'config_dir': settings.AIIDA_CONFIG_FOLDER, settings.DEFAULT_AIIDA_PATH_VARIABLE: os.environ.get(settings.DEFAULT_AIIDA_PATH_VARIABLE, None) }
@property def profile_dictionary(self): """Profile parameters. Used to set up AiiDA profile from self.profile_info dictionary. """ dictionary = { 'test_profile': True, 'storage': { 'backend': self.profile_info.get('storage_backend'), 'config': { 'database_engine': self.profile_info.get('database_engine'), 'database_port': self.profile_info.get('database_port'), 'database_hostname': self.profile_info.get('database_hostname'), 'database_name': self.profile_info.get('database_name'), 'database_username': self.profile_info.get('database_username'), 'database_password': self.profile_info.get('database_password'), 'repository_uri': f'file://{self.repo}', } }, 'process_control': { 'backend': 'rabbitmq', 'config': { 'broker_protocol': self.profile_info.get('broker_protocol'), 'broker_username': self.profile_info.get('broker_username'), 'broker_password': self.profile_info.get('broker_password'), 'broker_host': self.profile_info.get('broker_host'), 'broker_port': self.profile_info.get('broker_port'), 'broker_virtual_host': self.profile_info.get('broker_virtual_host'), } } } return dictionary
[docs] def create_db_cluster(self): """ Create the database cluster using PGTest. """ from pgtest.pgtest import PGTest if self.pg_cluster is not None: raise TestManagerError( 'Running temporary postgresql cluster detected.Use destroy_all() before creating a new cluster.' ) self.pg_cluster = PGTest(**self._pgtest) self.dbinfo.update(self.pg_cluster.dsn)
[docs] def create_aiida_db(self): """ Create the necessary database on the temporary postgres instance. """ if configuration.get_profile() is not None: raise TestManagerError('An AiiDA profile can not be loaded while creating a tests db environment') if self.pg_cluster is None: self.create_db_cluster() self.postgres = Postgres(interactive=False, quiet=True, dbinfo=self.dbinfo) # Note: We give the user CREATEDB privileges here, only since they are required for the migration tests self.postgres.create_dbuser( self.profile_info['database_username'], self.profile_info['database_password'], 'CREATEDB' ) self.postgres.create_db(self.profile_info['database_username'], self.profile_info['database_name']) self.dbinfo = self.postgres.dbinfo self.profile_info['database_hostname'] = self.postgres.host_for_psycopg2 self.profile_info['database_port'] = self.postgres.port_for_psycopg2 self._has_test_db = True
[docs] def create_profile(self): """ Set AiiDA to use the tests config dir and create a default profile there Warning: the AiiDA dbenv must not be loaded when this is called! """ from aiida.manage.configuration import Profile from aiida.orm import User manager = get_manager() if not self._has_test_db: self.create_aiida_db() if not self.root_dir: self.root_dir = tempfile.TemporaryDirectory().name # pylint: disable=consider-using-with configuration.CONFIG = None os.environ[settings.DEFAULT_AIIDA_PATH_VARIABLE] = self.config_dir settings.set_configuration_directory() manager.unload_profile() profile_name = self.profile_info['name'] config = configuration.get_config(create=True) profile = Profile(profile_name, self.profile_dictionary) config.add_profile(profile) config.set_default_profile(profile_name).store() self._profile = profile # initialise the profile profile = manager.load_profile(profile_name) # initialize the profile storage profile.storage_cls.migrate(profile) # create the default user for the profile created, user = User.collection.get_or_create(**get_user_dict(_DEFAULT_PROFILE_INFO)) if created: user.store() profile.default_user_email = user.email
[docs] def repo_ok(self): return bool(self.repo and os.path.isdir(os.path.dirname(self.repo)))
@property def repo(self): return self._return_dir(self.profile_info['repo_dir'])
[docs] def _return_dir(self, dir_path): """Return a path to a directory from the fs environment""" if os.path.isabs(dir_path): return dir_path return os.path.join(self.root_dir, dir_path)
@property def backend(self): return self.profile_info['backend'] @backend.setter def backend(self, backend): if self.has_profile_open(): raise TestManagerError('backend cannot be changed after setting up the environment') valid_backends = ['core.psql_dos'] if backend not in valid_backends: raise ValueError(f'invalid backend {backend}, must be one of {valid_backends}') self.profile_info['backend'] = backend @property def config_dir_ok(self): return bool(self.config_dir and os.path.isdir(self.config_dir)) @property def config_dir(self): return self._return_dir(self.profile_info['config_dir']) @property def root_dir(self): return self.profile_info['root_path'] @root_dir.setter def root_dir(self, root_dir): self.profile_info['root_path'] = root_dir @property def root_dir_ok(self): return bool(self.root_dir and os.path.isdir(self.root_dir))
[docs] def destroy_all(self): """Remove all traces of the tests run""" super().destroy_all() if self.root_dir: shutil.rmtree(self.root_dir) self.root_dir = None if self.pg_cluster: self.pg_cluster.close() self.pg_cluster = None self._has_test_db = False self._profile = None if 'config' in self._backup: configuration.CONFIG = self._backup['config'] if 'config_dir' in self._backup: settings.AIIDA_CONFIG_FOLDER = self._backup['config_dir'] if settings.DEFAULT_AIIDA_PATH_VARIABLE in self._backup and self._backup[settings.DEFAULT_AIIDA_PATH_VARIABLE]: os.environ[settings.DEFAULT_AIIDA_PATH_VARIABLE] = self._backup[settings.DEFAULT_AIIDA_PATH_VARIABLE]
[docs] def has_profile_open(self): return self._profile is not None
_GLOBAL_TEST_MANAGER = TestManager()
[docs]@contextmanager def test_manager(backend='core.psql_dos', profile_name=None, pgtest=None): """ Context manager for TestManager objects. Sets up temporary AiiDA environment for testing or reuses existing environment, if `AIIDA_TEST_PROFILE` environment variable is set. Example pytest fixture:: def aiida_profile(): with test_manager(backend) as test_mgr: yield fixture_mgr Example unittest test runner:: with test_manager(backend) as test_mgr: # ready for tests # everything cleaned up :param backend: storage backend type name :param profile_name: name of test profile to be used or None (to use temporary profile) :param pgtest: a dictionary of arguments to be passed to PGTest() for starting the postgresql cluster, e.g. {'pg_ctl': '/somepath/pg_ctl'}. Should usually not be necessary. """ from aiida.common.log import configure_logging from aiida.common.utils import Capturing try: if not _GLOBAL_TEST_MANAGER.has_profile_open(): if profile_name: _GLOBAL_TEST_MANAGER.use_profile(profile_name=profile_name) else: with Capturing(): # capture output of AiiDA DB setup _GLOBAL_TEST_MANAGER.use_temporary_profile(backend=backend, pgtest=pgtest) configure_logging(with_orm=True) yield _GLOBAL_TEST_MANAGER finally: _GLOBAL_TEST_MANAGER.destroy_all()
[docs]def get_test_backend_name() -> str: """ Read name of storage backend from environment variable or the specified test profile. Reads storage backend from 'AIIDA_TEST_BACKEND' environment variable, or the backend configured for the 'AIIDA_TEST_PROFILE'. :returns: name of storage backend :raises: ValueError if unknown backend name detected. :raises: ValueError if both 'AIIDA_TEST_BACKEND' and 'AIIDA_TEST_PROFILE' are set, and the two backends do not match. """ test_profile_name = get_test_profile_name() backend_env = os.environ.get('AIIDA_TEST_BACKEND', None) if test_profile_name is not None: backend_profile = configuration.get_config().get_profile(test_profile_name).storage_backend if backend_env is not None and backend_env != backend_profile: raise ValueError( "The backend '{}' read from AIIDA_TEST_BACKEND does not match the backend '{}' " "of AIIDA_TEST_PROFILE '{}'".format(backend_env, backend_profile, test_profile_name) ) backend_res = backend_profile else: backend_res = backend_env or 'core.psql_dos' if backend_res in ('core.psql_dos',): return backend_res raise ValueError(f"Unknown backend '{backend_res}' read from AIIDA_TEST_BACKEND environment variable")
[docs]def get_test_profile_name(): """ Read name of test profile from environment variable. Reads name of existing test profile 'AIIDA_TEST_PROFILE' environment variable. If specified, this profile is used for running the tests (instead of setting up a temporary profile). :returns: content of environment variable or `None` """ return os.environ.get('AIIDA_TEST_PROFILE', None)
[docs]def get_user_dict(profile_dict): """Collect parameters required for creating users.""" return {k: profile_dict[k] for k in ('email', 'first_name', 'last_name', 'institution')}