Allow to use Engine subclass in tests

README.md

@@ -82,6 +82,24 @@ async def test_returns_true_too(db):
         assert await con.fetchval("SELECT true")
 ```
 
+If you want to use in tests your own custom `Engine` subclass you can define `asyncpg_engine_cls` fixture that returns it:
+```python
+from asyncpg_engine import Engine
+
+
+class MyPrettyEngine(Engine):
+    pass
+
+
+@pytest.fixture()
+def asyncpg_engine_cls() -> t.Type[MyPrettyEngine]:
+    return MyPrettyEngine
+
+
+async def test_returns_my_pretty_engine(db: MyPrettyEngine) -> None:
+    assert isinstance(db, MyPrettyEngine)
+```
+
 ## Development and contribution
 
 First of all you should install Poetry using [official instructions](https://python-poetry.org/docs/#osx--linux--bashonwindows-install-instructions) or solutions provided by your distro. Then install dependencies:

asyncpg_engine/engine.py

@@ -21,9 +21,7 @@ def __init__(self, pool: Pool, use_single_connection: bool):
         self._global_con = None
 
     @classmethod
-    async def create(
-        cls, url: str, *, use_single_connection: bool = False, **kwargs: Any
-    ) -> "Engine":
+    async def create(cls, url: str, *, use_single_connection: bool = False, **kwargs: Any) -> "Engine":
         pool = await create_pool(url, min_size=2, init=cls._set_codecs, **kwargs)
         return cls(pool, use_single_connection)
 
@@ -74,7 +72,5 @@ def __await__(self) -> Generator[Connection, None, None]:
     async def __aenter__(self) -> Connection:
         return await self()
 
-    async def __aexit__(
-        self, exc_type: Type[BaseException], exc: BaseException, tv: TracebackType
-    ) -> None:
+    async def __aexit__(self, exc_type: Type[BaseException], exc: BaseException, tv: TracebackType) -> None:
         await self.engine.release(self.con)

asyncpg_engine/pytest_plugin.py

@@ -7,20 +7,23 @@
 
 
 def pytest_configure(config: t.Any) -> None:
-    config.addinivalue_line(
-        "markers", "asyncpg_engine: configure asyncpg-engine plugin behaviour."
-    )
+    config.addinivalue_line("markers", "asyncpg_engine: configure asyncpg-engine plugin behaviour.")
+
+
+@pytest.fixture()
+def asyncpg_engine_cls() -> t.Type[Engine]:
+    return Engine
 
 
 @pytest.fixture()
 async def db(
-    request: pytest.FixtureRequest, postgres_url: str
+    request: pytest.FixtureRequest, asyncpg_engine_cls: t.Type[Engine], postgres_url: str
 ) -> t.AsyncGenerator[Engine, None]:
     plugin_config = request.node.get_closest_marker("asyncpg_engine")
 
     transactional = getattr(plugin_config, "kwargs", {}).get("transactional", True)
 
-    _db = await Engine.create(url=postgres_url, use_single_connection=transactional)
+    _db = await asyncpg_engine_cls.create(url=postgres_url, use_single_connection=transactional)
 
     con: Connection = await _db.acquire()
     tr = con.transaction()

pyproject.toml

@@ -58,3 +58,6 @@ warn_redundant_casts = true
 warn_return_any = true
 warn_unused_ignores = true
 implicit_reexport = false
+
+[tool.black]
+line-length = 120

tests/pytest_plugin/test_custom_engine_cls.py

@@ -0,0 +1,20 @@
+import typing as t
+
+import pytest
+
+from asyncpg_engine import Engine
+
+pytestmark = [pytest.mark.asyncio]
+
+
+class MyPrettyEngine(Engine):
+    pass
+
+
+@pytest.fixture()
+def asyncpg_engine_cls() -> t.Type[MyPrettyEngine]:
+    return MyPrettyEngine
+
+
+async def test_returns_my_pretty_engine(db: MyPrettyEngine) -> None:
+    assert isinstance(db, MyPrettyEngine)

tests/pytest_plugin/test_non_transactional.py

@@ -0,0 +1,27 @@
+import pytest
+from asyncpg import Connection
+
+from asyncpg_engine import Engine
+
+pytestmark = [pytest.mark.asyncio, pytest.mark.asyncpg_engine(transactional=False)]
+
+
+async def test_acquire_returns_new_connection(db: Engine) -> None:
+    async with db.acquire() as con_0:
+        async with db.acquire() as con_1:
+
+            assert con_0 != con_1
+
+
+async def test_con_is_not_the_same_as_acquire_result(db: Engine, con: Connection) -> None:
+    async with db.acquire() as con_0:
+        async with db.acquire() as con_1:
+
+            assert con_0 != con_1
+
+
+@pytest.mark.asyncpg_engine(transactional=False)
+async def test_not_transactional(con: Connection) -> None:
+    sql = "SELECT txid_current()"
+
+    assert await con.fetchval(sql) != await con.fetchval(sql)

tests/pytest_plugin/test_transactional.py

@@ -0,0 +1,32 @@
+import pytest
+from asyncpg import Connection
+
+from asyncpg_engine import Engine
+
+pytestmark = [pytest.mark.asyncio]
+
+
+async def test_acquire_returns_the_same_connection(db: Engine) -> None:
+    async with db.acquire() as con_0:
+        async with db.acquire() as con_1:
+
+            assert con_0 == con_1
+
+
+async def test_con_is_the_same_as_acquire_result(db: Engine, con: Connection) -> None:
+    async with db.acquire() as con_0:
+
+        assert con_0 == con
+
+
+async def test_transactional(con: Connection) -> None:
+    sql = "SELECT txid_current()"
+
+    assert await con.fetchval(sql) == await con.fetchval(sql)
+
+
+@pytest.mark.asyncpg_engine(transactional=True)
+async def test_transactional_if_specified(con: Connection) -> None:
+    sql = "SELECT txid_current()"
+
+    assert await con.fetchval(sql) == await con.fetchval(sql)

tests/test_engine.py

@@ -18,9 +18,7 @@ async def test_returns_new_connection_every_acquire(engine: Engine) -> None:
             assert con_0 != con_1
 
 
-async def test_returns_the_same_connection_every_acquire_if_single(
-    postgres_url: str,
-) -> None:
+async def test_returns_the_same_connection_every_acquire_if_single(postgres_url: str) -> None:
     engine = await Engine.create(postgres_url, use_single_connection=True)
 
     async with engine.acquire() as con_0:
@@ -29,9 +27,7 @@ async def test_returns_the_same_connection_every_acquire_if_single(
             assert con_0 == con_1
 
 
-async def test_non_force_release_ignored_for_single_connection(
-    postgres_url: str,
-) -> None:
+async def test_non_force_release_ignored_for_single_connection(postgres_url: str) -> None:
     engine = await Engine.create(postgres_url, use_single_connection=True)
     async with engine.acquire() as con_0:
         pass