From a87b2344026f5ad9157ce81db1dc5e35a06d50ba Mon Sep 17 00:00:00 2001 From: Dan Cardin Date: Tue, 8 Feb 2022 15:00:04 -0500 Subject: [PATCH] feat: Add the ability to specify custom engine kwargs. --- pyproject.toml | 2 +- src/pytest_mock_resources/container/postgres.py | 16 ++++------------ .../fixture/database/relational/postgresql.py | 9 ++++++--- 3 files changed, 11 insertions(+), 16 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 0963a9c5..431dcb54 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "pytest-mock-resources" -version = "2.1.11" +version = "2.1.12" description = "A pytest plugin for easily instantiating reproducible mock resources." authors = [ "Omar Khan ", diff --git a/src/pytest_mock_resources/container/postgres.py b/src/pytest_mock_resources/container/postgres.py index 1ef1fa51..665fcb8b 100644 --- a/src/pytest_mock_resources/container/postgres.py +++ b/src/pytest_mock_resources/container/postgres.py @@ -49,10 +49,8 @@ def root_database(self): raise NotImplementedError() -def get_sqlalchemy_engine(config, database_name, isolation_level=None): - URI_TEMPLATE = ( - "postgresql+psycopg2://{username}:{password}@{host}:{port}/{database}?sslmode=disable" - ) +def get_sqlalchemy_engine(config, database_name, **engine_kwargs): + URI_TEMPLATE = "postgresql+psycopg2://{username}:{password}@{host}:{port}/{database}?sslmode=disable" DB_URI = URI_TEMPLATE.format( host=config.host, port=config.port, @@ -61,16 +59,12 @@ def get_sqlalchemy_engine(config, database_name, isolation_level=None): database=database_name, ) - options = {} - if isolation_level: - options["isolation_level"] = isolation_level - # Trigger any psycopg2-based import failures from pytest_mock_resources.compat import psycopg2 psycopg2.connect - engine = sqlalchemy.create_engine(DB_URI, **options) + engine = sqlalchemy.create_engine(DB_URI, **engine_kwargs) # Verify engine is connected engine.connect() @@ -83,9 +77,7 @@ def check_postgres_fn(config): get_sqlalchemy_engine(config, config.root_database) except sqlalchemy.exc.OperationalError: raise ContainerCheckFailed( - "Unable to connect to a presumed Postgres test container via given config: {}".format( - config - ) + "Unable to connect to a presumed Postgres test container via given config: {}".format(config) ) diff --git a/src/pytest_mock_resources/fixture/database/relational/postgresql.py b/src/pytest_mock_resources/fixture/database/relational/postgresql.py index 4fbe9d2e..f34dbeb3 100644 --- a/src/pytest_mock_resources/fixture/database/relational/postgresql.py +++ b/src/pytest_mock_resources/fixture/database/relational/postgresql.py @@ -19,11 +19,11 @@ def pmr_postgres_config(): def create_engine_manager( - pmr_postgres_config, ordered_actions, tables, createdb_template="template1" + pmr_postgres_config, ordered_actions, tables, createdb_template="template1", engine_kwargs=None ): database_name = produce_clean_database(pmr_postgres_config, createdb_template=createdb_template) - engine = get_sqlalchemy_engine(pmr_postgres_config, database_name) + engine = get_sqlalchemy_engine(pmr_postgres_config, database_name, **(engine_kwargs or {})) assign_fixture_credentials( engine, drivername="postgresql+psycopg2", @@ -42,7 +42,8 @@ def create_postgres_fixture( tables=None, session=None, async_=False, - createdb_template="template1" + createdb_template="template1", + engine_kwargs=None, ): """Produce a Postgres fixture. @@ -59,11 +60,13 @@ def create_postgres_fixture( 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_kwargs: Optional set of kwargs to send into the engine on creation. """ engine_manager_kwargs = dict( ordered_actions=ordered_actions, tables=tables, createdb_template=createdb_template, + engine_kwargs=engine_kwargs, ) @pytest.fixture(scope=scope)