# -*- 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
import warnings
from aiida.common.warnings import AiidaDeprecationWarning
from aiida.manage import configuration, get_manager
from aiida.manage.configuration.settings import create_instance_directories
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': '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):
warnings.warn('reset_db() is deprecated, use clear_profile() instead', AiidaDeprecationWarning)
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.reset_profile()
manager.get_profile_storage() # reload the storage connection
[docs] def has_profile_open(self):
return self._profile is not None
[docs] def destroy_all(self):
pass
[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='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.
"""
from aiida.manage.configuration import settings
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,
}
@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, settings
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.mkdtemp()
configuration.CONFIG = None
settings.AIIDA_CONFIG_FOLDER = self.config_dir
manager.unload_profile()
create_instance_directories()
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.objects.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 = ['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"""
from aiida.manage.configuration import settings
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']
[docs] def has_profile_open(self):
return self._profile is not None
_GLOBAL_TEST_MANAGER = TestManager()
[docs]@contextmanager
def test_manager(backend='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 'psql_dos'
if backend_res in ('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')}