# -*- 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.
"""
import contextlib
import os
import shutil
import tempfile
import warnings
from aiida.common.log import override_log_level
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.
If the daemon is running, it will be stopped because it might be holding on to entities that will be cleared
from the storage backend.
"""
from aiida.engine.daemon.client import get_daemon_client
daemon_client = get_daemon_client()
if daemon_client.is_daemon_running:
daemon_client.stop_daemon(wait=True)
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
with warnings.catch_warnings():
warnings.filterwarnings('ignore', category=UserWarning)
# This will raise a warning that the ``.aiida`` configuration directory is created.
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
# Load the new profile and initialize the profile storage
with override_log_level():
profile = manager.load_profile(profile_name)
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
# Set options to suppress certain warnings
config.set_option('warnings.development_version', False)
config.set_option('warnings.rabbitmq_version', False)
config.store()
[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]@contextlib.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')}