Source code for aiida.tools.pytest_fixtures.storage

"""Fixtures providing resources for storage plugins."""

from __future__ import annotations

import pathlib
import typing as t
from uuid import uuid4

import pytest
from pgtest.pgtest import PGTest

if t.TYPE_CHECKING:
    from pgtest.pgtest import PGTest


[docs] class PostgresCluster:
[docs] def __init__(self): # We initialize the cluster lazily self.cluster = None
[docs] def _create(self): try: self.cluster = PGTest() except OSError as e: raise RuntimeError('Could not initialize PostgreSQL cluster') from e
[docs] def _close(self): if self.cluster is not None: self.cluster.close()
[docs] def create_database( self, database_name: str | None = None, database_username: str | None = None, database_password: str | None = None, ) -> dict[str, str]: from aiida.manage.external.postgres import Postgres if self.cluster is None: self._create() postgres_config = { 'database_engine': 'postgresql_psycopg2', 'database_name': database_name or str(uuid4()), 'database_username': database_username or 'guest', 'database_password': database_password or 'guest', } postgres = Postgres(interactive=False, quiet=True, dbinfo=self.cluster.dsn) # type: ignore[union-attr] if not postgres.dbuser_exists(postgres_config['database_username']): postgres.create_dbuser( postgres_config['database_username'], postgres_config['database_password'], 'CREATEDB' ) postgres.create_db(postgres_config['database_username'], postgres_config['database_name']) postgres_config['database_hostname'] = postgres.host_for_psycopg2 postgres_config['database_port'] = postgres.port_for_psycopg2 return postgres_config
# TODO: Update docstring accordingly
[docs] @pytest.fixture(scope='session') def postgres_cluster(): """Create a temporary and isolated PostgreSQL cluster using ``pgtest`` and cleanup after the yield. :param database_name: Name of the database. :param database_username: Username to use for authentication. :param database_password: Password to use for authentication. :returns: Dictionary with parameters to connect to the PostgreSQL cluster. """ cluster = PostgresCluster() yield cluster cluster._close()
[docs] @pytest.fixture(scope='session') def config_psql_dos( tmp_path_factory: pytest.TempPathFactory, postgres_cluster: 'PGTest', ) -> t.Callable[[str | None, str | None, str | None], dict[str, t.Any]]: """Return a profile configuration for the :class:`~aiida.storage.psql_dos.backend.PsqlDosBackend`. The factory has the following signature to allow further configuring the database that is created: :param database_name: Name of the database to be created. :param database_username: Username to use for authentication. :param database_password: Password to use for authentication. :returns: The dictionary with the storage configuration for the ``core.psql_dos`` storage plugin. """ def factory( database_name: str | None = None, database_username: str | None = None, database_password: str | None = None ) -> dict[str, t.Any]: storage_config: dict[str, t.Any] = postgres_cluster.create_database( database_name=database_name, database_username=database_username, database_password=database_password, ) storage_config['repository_uri'] = f'file://{tmp_path_factory.mktemp("repository")}' return storage_config return factory
[docs] @pytest.fixture(scope='session') def config_sqlite_dos( tmp_path_factory: pytest.TempPathFactory, ) -> t.Callable[[str | pathlib.Path | None], dict[str, t.Any]]: """Return a profile configuration for the :class:`~aiida.storage.sqlite_dos.backend.SqliteDosStorage`. The factory has the following signature to allow further configuring the database that is created: :param filepath: Optional path to the sqlite database file. :returns: The dictionary with the storage configuration for the ``core.sqlite_dos`` storage plugin. """ def factory(filepath: str | pathlib.Path | None = None) -> dict[str, t.Any]: return {'filepath': str(filepath or tmp_path_factory.mktemp('test_sqlite_dos_storage'))} return factory