Merge pull request #135 from schireson/dc/configurable-createdb-template

Makefile

@@ -27,11 +27,11 @@ test: test-parallel
 
 ## Lint
 lint:
-	flake8 src tests
-	isort --check-only --diff src tests
-	pydocstyle src tests
-	black --check src tests
-	mypy src tests
+	flake8 src tests || exit 1
+	isort --check-only --diff src tests || exit 1
+	pydocstyle src tests || exit 1
+	black --check src tests || exit 1
+	mypy src tests || exit 1
 
 format:
 	isort --recursive src tests

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>",
@@ -60,6 +60,7 @@ pytest-asyncio = "^0.15.1"
 types-six = "^1.16.0"
 types-PyMySQL = "^1.0.2"
 types-redis = "^3.5.6"
+sqlalchemy2-stubs = "^0.0.2-alpha.19"
 
 [tool.poetry.extras]
 postgres = ['psycopg2']
@@ -87,6 +88,13 @@ use_parentheses = true
 [tool.black]
 line_length = 100
 
+[tool.mypy]
+strict_optional = true
+ignore_missing_imports = true
+warn_unused_ignores = true
+incremental = true
+plugins = 'sqlalchemy.ext.mypy.plugin'
+
 [build-system]
 requires = ["poetry_core==1.0.4"]
 build-backend = "poetry.core.masonry.api"

src/pytest_mock_resources/cli.py

@@ -2,7 +2,7 @@
 import enum
 import subprocess  # nosec
 
-from pytest_mock_resources import PostgresConfig, MysqlConfig, MongoConfig
+from pytest_mock_resources import MongoConfig, MysqlConfig, PostgresConfig
 
 postgres_image = PostgresConfig().image
 mysql_image = MysqlConfig().image

src/pytest_mock_resources/compat/sqlalchemy.py

@@ -3,9 +3,9 @@
 from pytest_mock_resources.compat.import_ import ImportAdaptor
 
 try:
-    from sqlalchemy.ext import asyncio  # type: ignore
+    from sqlalchemy.ext import asyncio
 except ImportError:
-    asyncio = ImportAdaptor(
+    asyncio = ImportAdaptor(  # type: ignore
         "SQLAlchemy",
         "SQLAlchemy >= 1.4",
         fail_message="Cannot use sqlalchemy async features with SQLAlchemy < 1.4.\n",

src/pytest_mock_resources/fixture/database/mongo.py

@@ -53,7 +53,9 @@ def _create_clean_database(config):
     new_database.command("createUser", db_id, pwd=password, roles=["dbOwner"])
 
     #  pass back an authenticated db connection
-    limited_client = pymongo.MongoClient(config.host, config.port, username=db_id, password=password, authSource=db_id)
+    limited_client = pymongo.MongoClient(
+        config.host, config.port, username=db_id, password=password, authSource=db_id
+    )
     limited_db = limited_client[db_id]
 
     assign_fixture_credentials(

src/pytest_mock_resources/fixture/database/redis.py

@@ -46,13 +46,15 @@ def create_redis_fixture(scope="function"):
     @pytest.fixture(scope=scope)
     def _(request, _redis_container, pmr_redis_config):
         database_number = 0
-        if hasattr(request.config, 'workerinput'):
+        if hasattr(request.config, "workerinput"):
             worker_input = request.config.workerinput
-            worker_id = worker_input['workerid']  # For example "gw0".
+            worker_id = worker_input["workerid"]  # For example "gw0".
             database_number = int(worker_id[2:])
 
         if database_number >= 16:
-            raise ValueError("The redis fixture currently only supports up to 16 parallel executions")
+            raise ValueError(
+                "The redis fixture currently only supports up to 16 parallel executions"
+            )
 
         db = redis.Redis(host=pmr_redis_config.host, port=pmr_redis_config.port, db=database_number)
         db.flushdb()

src/pytest_mock_resources/fixture/database/relational/postgresql.py

@@ -18,8 +18,11 @@ 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 = produce_clean_database(pmr_postgres_config, createdb_template=createdb_template)
+
     engine = get_sqlalchemy_engine(pmr_postgres_config, database_name)
     assign_fixture_credentials(
         engine,
@@ -34,7 +37,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 +57,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,11 +83,17 @@ async def _async(_postgres_container, pmr_postgres_config):
         return _sync
 
 
-def _create_clean_database(config):
+def produce_clean_database(config, createdb_template="template1"):
     root_engine = get_sqlalchemy_engine(config, config.root_database, isolation_level="AUTOCOMMIT")
+    with root_engine.connect() as conn:
+        database_name = _create_clean_database(conn, createdb_template=createdb_template)
+
+    return database_name
+
 
+def _create_clean_database(conn, createdb_template="template1"):
     try:
-        root_engine.execute(
+        conn.execute(
             """
             CREATE TABLE IF NOT EXISTS pytest_mock_resource_db(
                 id serial
@@ -87,15 +108,11 @@ def _create_clean_database(config):
         #  - the current process tries to commit the table creation
         pass
 
-    result = root_engine.execute(
-        "INSERT INTO pytest_mock_resource_db VALUES (DEFAULT) RETURNING id"
-    )
+    result = conn.execute("INSERT INTO pytest_mock_resource_db VALUES (DEFAULT) RETURNING id")
     id_ = tuple(result)[0][0]
     database_name = "pytest_mock_resource_db_{}".format(id_)
 
-    root_engine.execute('CREATE DATABASE "{}"'.format(database_name))
-    root_engine.execute(
-        'GRANT ALL PRIVILEGES ON DATABASE "{}" TO CURRENT_USER'.format(database_name)
-    )
+    conn.execute('CREATE DATABASE "{}" template={}'.format(database_name, createdb_template))
+    conn.execute('GRANT ALL PRIVILEGES ON DATABASE "{}" TO CURRENT_USER'.format(database_name))
 
     return database_name

src/pytest_mock_resources/fixture/database/relational/redshift/__init__.py

@@ -3,8 +3,8 @@
 from pytest_mock_resources.fixture.database.generic import assign_fixture_credentials
 from pytest_mock_resources.fixture.database.relational.generic import EngineManager
 from pytest_mock_resources.fixture.database.relational.postgresql import (
-    _create_clean_database,
     get_sqlalchemy_engine,
+    produce_clean_database,
 )
 from pytest_mock_resources.patch.redshift import psycopg2, sqlalchemy
 
@@ -30,7 +30,7 @@ def create_redshift_fixture(*ordered_actions, scope="function", tables=None, ses
 
     @pytest.fixture(scope=scope)
     def _(_redshift_container, pmr_postgres_config):
-        database_name = _create_clean_database(pmr_postgres_config)
+        database_name = produce_clean_database(pmr_postgres_config)
         engine = get_sqlalchemy_engine(pmr_postgres_config, database_name)
 
         assign_fixture_credentials(

src/pytest_mock_resources/fixture/database/relational/sqlite.py

@@ -18,7 +18,7 @@
 
 import pytest
 from sqlalchemy import create_engine, event
-from sqlalchemy.dialects import registry
+from sqlalchemy.dialects import registry  # type: ignore
 from sqlalchemy.dialects.postgresql import JSON, JSONB
 from sqlalchemy.dialects.sqlite import base as sqlite_base
 from sqlalchemy.dialects.sqlite.pysqlite import SQLiteDialect_pysqlite

tests/examples/test_multiprocess_redis_database/test_split.py

@@ -10,8 +10,6 @@
 A correct implementation would use some mechanism to avoid this inter-parallel-test
 key conflict problem.
 """
-import random
-import time
 
 
 def test_node_one(redis, pytestconfig):
@@ -31,8 +29,8 @@ def test_node_four(redis, pytestconfig):
 
 
 def run_test(redis, pytestconfig):
-    worker_id = int(pytestconfig.workerinput['workerid'][2:])
-    database = redis.connection_pool.get_connection('set').db
+    worker_id = int(pytestconfig.workerinput["workerid"][2:])
+    database = redis.connection_pool.get_connection("set").db
     assert worker_id == database
     print(worker_id, database)
 
@@ -41,7 +39,7 @@ def run_test(redis, pytestconfig):
     # XXX: however until the plugin is overall more process-safe, it's too flaky.
     # time.sleep(random.randrange(1, 10) / 10)
     value = redis.get("foo")
-    
+
     assert value == b"bar"
 
     redis.flushdb()

tests/fixture/database/__init__.py

@@ -1,9 +1,10 @@
 import random
 
+from sqlalchemy import create_engine
+
 from pytest_mock_resources.compat import boto3, moto, psycopg2
 from pytest_mock_resources.patch.redshift.mock_s3_copy import read_data_csv
 from pytest_mock_resources.patch.redshift.mock_s3_unload import get_data_csv
-from sqlalchemy import create_engine
 
 original_data = [
     (3342, 32434.0, "a", "gfhsdgaf"),

tests/fixture/database/test_ordered_actions.py

@@ -1,3 +1,5 @@
+from typing import List
+
 import pytest
 from sqlalchemy import Column, ForeignKey, Integer, String, text
 from sqlalchemy.exc import ProgrammingError
@@ -15,7 +17,8 @@ class User(Base):
 
     id = Column(Integer, primary_key=True, autoincrement=True)
     name = Column(String, nullable=False)
-    objects = relationship("Object", back_populates="owner")
+
+    objects: List["Object"] = relationship("Object", back_populates="owner")
 
 
 class Object(Base):
@@ -25,7 +28,8 @@ class Object(Base):
     id = Column(Integer, primary_key=True, autoincrement=True)
     name = Column(String, nullable=False)
     belongs_to = Column(Integer, ForeignKey("stuffs.user.id"))
-    owner = relationship("User", back_populates="objects")
+
+    owner: List[User] = relationship("User", back_populates="objects")
 
 
 rows = Rows(User(name="Harold"), User(name="Gump"))

tests/fixture/database/test_postgres.py

@@ -0,0 +1,46 @@
+from sqlalchemy import Column, event, Integer
+from sqlalchemy.ext.declarative import declarative_base
+
+from pytest_mock_resources import create_postgres_fixture
+from pytest_mock_resources.container.postgres import get_sqlalchemy_engine
+from pytest_mock_resources.fixture.database.relational.postgresql import _create_clean_database
+
+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")
+
+
+def test_create_clean_database_createdb_template(pmr_postgres_config, createdb_template_pg):
+    """Assert `createdb_template` is included in emitted CREATE DATABASE statement."""
+    root_engine = get_sqlalchemy_engine(
+        pmr_postgres_config, pmr_postgres_config.root_database, isolation_level="AUTOCOMMIT"
+    )
+
+    statement = ""
+
+    def before_execute(conn, clauseelement, multiparams, params, execution_options):
+        # Search for our create database statement, so we can assert against it.
+        if "CREATE DATABASE" in clauseelement:
+            nonlocal statement
+            statement = clauseelement
+        return clauseelement, multiparams, params
+
+    # Use the event system to hook into the statements being executed by sqlalchemy.
+    with root_engine.connect().execution_options(isolation_level="AUTOCOMMIT") as conn:
+        event.listen(conn, "before_execute", before_execute)
+        _create_clean_database(conn, createdb_template="template0")
+        event.remove(conn, "before_execute", before_execute)
+
+    assert "template0" in statement
+
+
+def test_createdb_template(createdb_template_pg):
+    """Assert successful usage of a fixture which sets the `createdb_template` argument."""
+    createdb_template_pg.execute(Thing.__table__.insert().values({"id": 1}))