From 35d7ca63b44f051a26d3f96d84e043919eb3f101 Mon Sep 17 00:00:00 2001
From: Sebastiaan Huber <mail@sphuber.net>
Date: Wed, 31 Jan 2024 19:27:55 +0100
Subject: [PATCH] Make `postgres_cluster` and `config_psql_dos` fixtures
 configurable (#6254)

The `postgres_cluster` and `config_psql_dos` fixtures are made available
publicly in order to help plugin packages with unit testing.
Particularly, these fixtures should make it easy to generate test
profiles using the `core.psql_dos` storage plugin, generating the
postgres database in a test cluster that is created on-the-fly.

The problem was that the interface did not actually allow to configure
anything, such as the database name, but this was hardcoded. This would
work fine in case it was used just once, but the fixtures could not be
used to create a second profile, as the database name would already
exist.

The `postgres_cluster` fixture is turned into a factory, allowing the
database name, username and password to be customized. The unbound
method `create_database` is added to the `PGTest` instance that it
returns to easily create a new database.

The `config_psql_dos` fixture now actually forwards the relevant part of
the `custom_configuration` to the `postgres_cluster.create_database`
call.

Finally, the `aiida_profile_factory` fixture factory now makes the
`custom_configuration` argument actually optional, as originally was
intended.
---
 src/aiida/manage/tests/pytest_fixtures.py | 54 +++++++++++++++--------
 1 file changed, 35 insertions(+), 19 deletions(-)

diff --git a/src/aiida/manage/tests/pytest_fixtures.py b/src/aiida/manage/tests/pytest_fixtures.py
index d9c5fcd8d3..3976bfbde2 100644
--- a/src/aiida/manage/tests/pytest_fixtures.py
+++ b/src/aiida/manage/tests/pytest_fixtures.py
@@ -70,9 +70,7 @@ def aiida_caplog(caplog):
 
 
 @pytest.fixture(scope='session')
-def postgres_cluster(
-    database_name: str | None = None, database_username: str | None = None, database_password: str | None = None
-) -> t.Generator[dict[str, str], None, None]:
+def postgres_cluster():
     """Create a temporary and isolated PostgreSQL cluster using ``pgtest`` and cleanup after the yield.
 
     :param database_name: Name of the database.
@@ -82,27 +80,35 @@ def postgres_cluster(
     """
     from pgtest.pgtest import PGTest
 
-    from aiida.manage.external.postgres import Postgres
+    def create_database(
+        database_name: str | None = None, database_username: str | None = None, database_password: str | None = None
+    ) -> t.Generator[dict[str, str], None, None]:
+        from aiida.manage.external.postgres import Postgres
 
-    postgres_config = {
-        'database_engine': 'postgresql_psycopg2',
-        'database_name': database_name or str(uuid.uuid4()),
-        'database_username': database_username or 'guest',
-        'database_password': database_password or 'guest',
-    }
-
-    cluster = None
-    try:
-        cluster = PGTest()
+        postgres_config = {
+            'database_engine': 'postgresql_psycopg2',
+            'database_name': database_name or str(uuid.uuid4()),
+            'database_username': database_username or 'guest',
+            'database_password': database_password or 'guest',
+        }
 
         postgres = Postgres(interactive=False, quiet=True, dbinfo=cluster.dsn)
-        postgres.create_dbuser(postgres_config['database_username'], postgres_config['database_password'], 'CREATEDB')
+        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
 
-        yield postgres_config
+        return postgres_config
+
+    cluster = None
+    try:
+        cluster = PGTest()
+        cluster.create_database = create_database
+        yield cluster
     finally:
         if cluster is not None:
             cluster.close()
@@ -194,16 +200,26 @@ def factory(custom_configuration: dict[str, t.Any] | None = None) -> dict[str, t
         :param custom_configuration: Custom configuration to override default profile configuration.
         :returns: The profile configuration.
         """
+        custom_configuration = custom_configuration or {}
+        custom_configuration.setdefault('storage', {}).setdefault('config', {})
+        custom_storage_config = custom_configuration['storage']['config']
+
+        storage_config = postgres_cluster.create_database(
+            database_name=custom_storage_config.get('database_name'),
+            database_username=custom_storage_config.get('database_username'),
+            database_password=custom_storage_config.get('database_password'),
+        )
+        recursive_merge(custom_storage_config, storage_config)
+
         configuration = {
             'storage': {
                 'backend': 'core.psql_dos',
                 'config': {
-                    **postgres_cluster,
                     'repository_uri': f'file://{tmp_path_factory.mktemp("repository")}',
                 },
             }
         }
-        recursive_merge(configuration, custom_configuration or {})
+        recursive_merge(configuration, custom_configuration)
         return configuration
 
     return factory
@@ -242,7 +258,7 @@ def aiida_profile_factory(
     contains and recreate the default user, making the profile ready for use.
     """
 
-    def factory(custom_configuration: dict[str, t.Any]) -> Profile:
+    def factory(custom_configuration: dict[str, t.Any] | None = None) -> Profile:
         """Create an isolated AiiDA instance with a temporary and fully loaded profile.
 
         :param custom_configuration: Custom configuration to override default profile configuration.