Skip to content

Commit

Permalink
fix: Ensure connection gets closed and is compatible with asyncpg.
Browse files Browse the repository at this point in the history
  • Loading branch information
DanCardin committed Feb 24, 2023
1 parent ed488f5 commit 56cab7b
Show file tree
Hide file tree
Showing 3 changed files with 9 additions and 3 deletions.
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[tool.poetry]
name = "pytest-mock-resources"
version = "2.6.8"
version = "2.6.9"
description = "A pytest plugin for easily instantiating reproducible mock resources."
authors = [
"Omar Khan <oakhan3@gmail.com>",
Expand Down
2 changes: 1 addition & 1 deletion src/pytest_mock_resources/container/postgres.py
Original file line number Diff line number Diff line change
Expand Up @@ -107,7 +107,7 @@ def get_sqlalchemy_engine(
database=database_name,
)

if getattr(url.get_dialect(), "is_async"):
if getattr(url.get_dialect(), "is_async", None):
from sqlalchemy.ext.asyncio import create_async_engine

engine = create_async_engine(url, **engine_kwargs, isolation_level=isolation_level)
Expand Down
8 changes: 7 additions & 1 deletion src/pytest_mock_resources/fixture/postgresql.py
Original file line number Diff line number Diff line change
Expand Up @@ -109,14 +109,18 @@ async def _async(*_, pmr_postgres_container, pmr_postgres_config):

def _sync_fixture(pmr_config, engine_manager_kwargs, engine_kwargs):
root_engine = cast(Engine, get_sqlalchemy_engine(pmr_config, pmr_config.root_database))
retry(root_engine.connect, retries=DEFAULT_RETRIES)
conn = retry(root_engine.connect, retries=DEFAULT_RETRIES)
conn.close()
root_engine.dispose()

root_engine = cast(Engine, get_sqlalchemy_engine(pmr_config, pmr_config.root_database))
with root_engine.connect() as root_conn:
with root_conn.begin() as trans:
template_database, template_manager, engine_manager = create_engine_manager(
root_conn, **engine_manager_kwargs
)
trans.commit()
root_engine.dispose()

if template_manager:
assert template_database
Expand All @@ -132,10 +136,12 @@ def _sync_fixture(pmr_config, engine_manager_kwargs, engine_kwargs):

# Everything below is normal per-test context. We create a brand new database/engine/manager
# distinct from what might have been used for the template database.
root_engine = cast(Engine, get_sqlalchemy_engine(pmr_config, pmr_config.root_database))
with root_engine.connect() as root_conn:
with root_conn.begin() as trans:
database_name = _produce_clean_database(root_conn, createdb_template=template_database)
trans.commit()
root_engine.dispose()

engine = get_sqlalchemy_engine(pmr_config, database_name, **engine_kwargs)
yield from engine_manager.manage_sync(engine)
Expand Down

0 comments on commit 56cab7b

Please sign in to comment.