"""Fixtures to create and load temporary AiiDA configuration directories and profiles."""
from __future__ import annotations
import contextlib
import os
import pathlib
import secrets
import typing as t
import pytest
if t.TYPE_CHECKING:
from aiida.manage.configuration.config import Config
[docs]
@pytest.fixture(scope='session')
def aiida_config_factory():
"""Return a factory to create and load a new temporary AiiDA configuration directory.
The factory is a context manager that returns a loaded :class:`aiida.manage.configuration.config.Config`. It
requires a path on the local file system where the configuration directory is to be created as an argument. If
another configuration directory was already loaded that is automatically restored at the end of the context manager.
This way, any changes made to the configuration during the context are fully temporary and automatically undone
after the test.
Usage::
def test(aiida_config_factory, tmp_path_factory):
import secrets
with aiida_config_factory(tmp_path_factory.mktemp(secrets.token_hex(16))) as config:
yield config
The factory has the following signature to allow further configuring the profile that is created and loaded:
:param dirpath: The path to create the configuration directory in.
:returns `~aiida.manage.configuration.config.Config`: The loaded temporary config.
"""
@contextlib.contextmanager
def factory(dirpath: pathlib.Path):
from aiida.common.exceptions import MissingConfigurationError
from aiida.manage.configuration import get_config, reset_config, settings
try:
current_config = get_config()
except MissingConfigurationError:
current_config = None
current_path_variable = os.environ.get(settings.DEFAULT_AIIDA_PATH_VARIABLE)
reset_config()
dirpath_config = dirpath / settings.DEFAULT_CONFIG_DIR_NAME
os.environ[settings.DEFAULT_AIIDA_PATH_VARIABLE] = str(dirpath_config)
settings.set_configuration_directory(dirpath_config)
config = get_config(create=True)
try:
yield config
finally:
if current_config:
reset_config()
settings.set_configuration_directory(pathlib.Path(current_config.dirpath))
get_config()
if current_path_variable is None:
os.environ.pop(settings.DEFAULT_AIIDA_PATH_VARIABLE, None)
else:
os.environ[settings.DEFAULT_AIIDA_PATH_VARIABLE] = current_path_variable
return factory
[docs]
@pytest.fixture(scope='session')
def aiida_profile_factory():
"""Return a factory to create and load a new temporary AiiDA profile.
The factory is a context manager that returns a loaded :class:`aiida.manage.configuration.profile.Profile`. It
requires a :class:`aiida.manage.configuration.config.Config` instance to which the profile is added. If another
profile was already loaded, that is automatically restored at the end of the context manager. This way, any changes
made to the profile during the context are fully temporary and automatically undone after the test. The created
``Profile`` instance dynamically has the method ``reset_storage`` added which, when called, deletes all content of
the storage, recreating the default user. The daemon is also stopped if it was running.
Usage::
def test(aiida_config_tmp, aiida_profile_factory):
with aiida_profile_factory(aiida_config_tmp) as profile:
yield profile
The factory has the following signature to allow further configuring the profile that is created and loaded:
:param storage_backend: The storage plugin to use. Defaults to ``core.sqlite_dos``.
:param storage_config: The configuration to use for the selected storage plugin.
:param broker_backend: The broker plugin to use. Defaults to defining no broker.
:param broker_config: The configuration to use for the selected broker plugin.
:param name: The name of the profile. Defaults to a random string.
:param name: The email to use for the default user. Defaults to ``test@localhost``.
:returns `~aiida.manage.configuration.profile.Profile`: The loaded temporary profile.
"""
@contextlib.contextmanager
def factory(
config: 'Config',
*,
storage_backend: str = 'core.sqlite_dos',
storage_config: dict[str, t.Any] | None = None,
broker_backend: str | None = None,
broker_config: dict[str, t.Any] | None = None,
name: str | None = None,
email: str = 'test@localhost',
):
from aiida.manage.configuration import create_profile, profile_context
from aiida.manage.manager import get_manager
manager = get_manager()
storage_config = storage_config or {'filepath': str(pathlib.Path(config.dirpath) / 'storage')}
if broker_backend and broker_config is None:
broker_config = {
'broker_protocol': 'amqp',
'broker_username': 'guest',
'broker_password': 'guest',
'broker_host': '127.0.0.1',
'broker_port': 5672,
'broker_virtual_host': '',
}
profile = create_profile(
config,
storage_backend=storage_backend,
storage_config=storage_config,
broker_backend=broker_backend,
broker_config=broker_config,
name=name or secrets.token_hex(16),
email=email,
is_test_profile=True,
)
config.set_default_profile(profile.name)
def reset_storage():
"""Reset the storage of the profile.
This ensures that the contents of the profile are reset as well as the ``Manager``, which may hold
references to data that will be destroyed. The daemon will also be stopped if it was running.
"""
from aiida.engine.daemon.client import DaemonException, get_daemon_client
from aiida.orm import User
if broker_backend:
daemon_client = get_daemon_client()
if daemon_client.is_daemon_running:
try:
daemon_client.stop_daemon(wait=True)
except DaemonException:
pass
manager.get_profile_storage()._clear()
manager.reset_profile()
User(email=profile.default_user_email or email).store()
# Add the ``reset_storage`` method, such that users can empty the storage through the ``Profile`` instance that
# is returned by this fixture.
setattr(profile, 'reset_storage', reset_storage)
with profile_context(profile, allow_switch=True):
yield profile
return factory
[docs]
@pytest.fixture(scope='session', autouse=True)
def aiida_config(tmp_path_factory, aiida_config_factory):
"""Return a loaded temporary AiiDA configuration directory.
This fixture is session-scoped and used automatically as soon as these fixtures are imported.
:returns :class:`~aiida.manage.configuration.config.Config`: The loaded temporary config.
"""
with aiida_config_factory(tmp_path_factory.mktemp(secrets.token_hex(16))) as config:
yield config
[docs]
@pytest.fixture(scope='session', autouse=True)
def aiida_profile(aiida_config, aiida_profile_factory):
"""Return a loaded temporary AiiDA profile.
This fixture is session-scoped and used automatically as soon as these fixtures are imported. The profile defines
no broker and uses the ``core.sqlite_dos`` storage backend, meaning it requires no services to run.
:returns :class:`~aiida.manage.configuration.profile.Profile`: The loaded temporary profile.
"""
with aiida_profile_factory(aiida_config) as profile:
yield profile
[docs]
@pytest.fixture(scope='function')
def aiida_profile_clean(aiida_profile):
"""Return a loaded temporary AiiDA profile where the data storage is cleaned before the start of the test.
This is a function-scoped version of the ``aiida_profile`` fixture.
:returns :class:`~aiida.manage.configuration.profile.Profile`: The loaded temporary profile.
"""
aiida_profile.reset_storage()
yield aiida_profile
[docs]
@pytest.fixture(scope='class')
def aiida_profile_clean_class(aiida_profile):
"""Return a loaded temporary AiiDA profile where the data storage is cleaned before the start of the test.
This is a class-scoped version of the ``aiida_profile`` fixture.
:returns `~aiida.manage.configuration.profile.Profile`: The loaded temporary profile.
"""
aiida_profile.reset_storage()
yield aiida_profile
[docs]
@pytest.fixture(scope='function')
def aiida_config_tmp(tmp_path, aiida_config_factory):
"""Create and load a temporary AiiDA configuration directory.
This fixture is function-scoped and automatically restores any previously loaded config after the test.
:returns :class:`~aiida.manage.configuration.config.Config`: The loaded temporary config.
"""
with aiida_config_factory(tmp_path) as config:
yield config
[docs]
@pytest.fixture(scope='function')
def aiida_profile_tmp(aiida_config_tmp, aiida_profile_factory):
"""Create and load a temporary AiiDA profile.
This fixture is function-scoped and automatically restores any previously loaded profile after the test.
:returns :class:`~aiida.manage.configuration.profile.Profile`: The loaded temporary profile.
"""
with aiida_profile_factory(aiida_config_tmp) as profile:
yield profile
