From 6305307fda52c7e4e16d0ed7772034685430056b Mon Sep 17 00:00:00 2001
From: Dan Cardin <ddcardin@gmail.com>
Date: Fri, 7 Jan 2022 10:23:27 -0500
Subject: [PATCH] feat: Add configurable template option for postgres database
 creation.

---
 pyproject.toml                                |  2 +-
 .../fixture/database/relational/postgresql.py | 28 ++++++++++++++-----
 tests/fixture/database/test_postgres.py       | 22 +++++++++++++++
 3 files changed, 44 insertions(+), 8 deletions(-)
 create mode 100644 tests/fixture/database/test_postgres.py

diff --git a/pyproject.toml b/pyproject.toml
index adf26c97..da171e64 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pytest-mock-resources"
-version = "2.1.9"
+version = "2.1.10"
 description = "A pytest plugin for easily instantiating reproducible mock resources."
 authors = [
     "Omar Khan <oakhan3@gmail.com>",
diff --git a/src/pytest_mock_resources/fixture/database/relational/postgresql.py b/src/pytest_mock_resources/fixture/database/relational/postgresql.py
index acf3c84d..4fd3d3a2 100644
--- a/src/pytest_mock_resources/fixture/database/relational/postgresql.py
+++ b/src/pytest_mock_resources/fixture/database/relational/postgresql.py
@@ -18,8 +18,10 @@ def pmr_postgres_config():
     return PostgresConfig()
 
 
-def create_engine_manager(pmr_postgres_config, ordered_actions, tables):
-    database_name = _create_clean_database(pmr_postgres_config)
+def create_engine_manager(
+    pmr_postgres_config, ordered_actions, tables, createdb_template="template1"
+):
+    database_name = _create_clean_database(pmr_postgres_config, createdb_template=createdb_template)
     engine = get_sqlalchemy_engine(pmr_postgres_config, database_name)
     assign_fixture_credentials(
         engine,
@@ -34,7 +36,12 @@ def create_engine_manager(pmr_postgres_config, ordered_actions, tables):
 
 
 def create_postgres_fixture(
-    *ordered_actions, scope="function", tables=None, session=None, async_=False
+    *ordered_actions,
+    scope="function",
+    tables=None,
+    session=None,
+    async_=False,
+    createdb_template="template1"
 ):
     """Produce a Postgres fixture.
 
@@ -49,16 +56,23 @@ def create_postgres_fixture(
         session: Whether to return a session instead of an engine directly. This can
             either be a bool or a callable capable of producing a session.
         async_: Whether to return an async fixture/client.
+        createdb_template: The template database used to create sub-databases. "template1" is the
+            default chosen when no template is specified.
     """
+    engine_manager_kwargs = dict(
+        ordered_actions=ordered_actions,
+        tables=tables,
+        createdb_template=createdb_template,
+    )
 
     @pytest.fixture(scope=scope)
     def _sync(_postgres_container, pmr_postgres_config):
-        engine_manager = create_engine_manager(pmr_postgres_config, ordered_actions, tables)
+        engine_manager = create_engine_manager(pmr_postgres_config, **engine_manager_kwargs)
         yield from engine_manager.manage_sync(session=session)
 
     @pytest.fixture(scope=scope)
     async def _async(_postgres_container, pmr_postgres_config):
-        engine_manager = create_engine_manager(pmr_postgres_config, ordered_actions, tables)
+        engine_manager = create_engine_manager(pmr_postgres_config, **engine_manager_kwargs)
         async for engine in engine_manager.manage_async(session=session):
             yield engine
 
@@ -68,7 +82,7 @@ async def _async(_postgres_container, pmr_postgres_config):
         return _sync
 
 
-def _create_clean_database(config):
+def _create_clean_database(config, createdb_template="template1"):
     root_engine = get_sqlalchemy_engine(config, config.root_database, isolation_level="AUTOCOMMIT")
 
     try:
@@ -93,7 +107,7 @@ def _create_clean_database(config):
     id_ = tuple(result)[0][0]
     database_name = "pytest_mock_resource_db_{}".format(id_)
 
-    root_engine.execute('CREATE DATABASE "{}"'.format(database_name))
+    root_engine.execute('CREATE DATABASE "{}" template={}'.format(database_name, createdb_template))
     root_engine.execute(
         'GRANT ALL PRIVILEGES ON DATABASE "{}" TO CURRENT_USER'.format(database_name)
     )
diff --git a/tests/fixture/database/test_postgres.py b/tests/fixture/database/test_postgres.py
new file mode 100644
index 00000000..40680381
--- /dev/null
+++ b/tests/fixture/database/test_postgres.py
@@ -0,0 +1,22 @@
+from sqlalchemy import Column, Integer
+from sqlalchemy.ext.declarative import declarative_base
+
+from pytest_mock_resources import create_postgres_fixture
+
+Base = declarative_base()
+
+
+class Thing(Base):
+    __tablename__ = "thing"
+
+    id = Column(Integer, autoincrement=True, primary_key=True)
+
+
+createdb_template_pg = create_postgres_fixture(Base, createdb_template="template0", session=True)
+
+
+def test_createdb_template(createdb_template_pg):
+    """Assert successful usage of a fixture which sets the `createdb_template` argument."""
+    thing = Thing(id=1)
+    createdb_template_pg.add(thing)
+    createdb_template_pg.commit()