@@ -282,14 +282,15 @@ Sample below is simplified session fixture from
282282
283283.. code-block :: python
284284
285+ from psycopg import Connection
285286 from sqlalchemy import create_engine
286- from sqlalchemy.orm import scoped_session, sessionmaker
287+ from sqlalchemy.orm import scoped_session, sessionmaker, Session
287288 from sqlalchemy.pool import NullPool
288289 from zope.sqlalchemy import register
289290
290291
291292 @pytest.fixture
292- def db_session (postgresql ) :
293+ def db_session (postgresql : Connection) -> Iterator[Session] :
293294 """ Session for SQLAlchemy."""
294295 from pyramid_fullauth.models import Base
295296
@@ -307,7 +308,7 @@ Sample below is simplified session fixture from
307308
308309
309310 @pytest.fixture
310- def user (db_session ) :
311+ def user (db_session : Session) -> User :
311312 """ Test user fixture."""
312313 from pyramid_fullauth.models import User
313314 from tests.tools import DEFAULT_USER
@@ -318,7 +319,7 @@ Sample below is simplified session fixture from
318319 return new_user
319320
320321
321- def test_remove_last_admin (db_session , user ) :
322+ def test_remove_last_admin (db_session : pyramid_basemodel.Session , user : User) -> None :
322323 """
323324 Sample test checks internal login, but shows usage in tests with SQLAlchemy
324325 """
@@ -346,11 +347,14 @@ For this import DatabaseJanitor and use its init and drop methods:
346347
347348.. code-block :: python
348349
350+ import psycopg
351+ from psycopg import Connection
349352 import pytest
350353 from pytest_postgresql.janitor import DatabaseJanitor
354+ from pytest_postgresql.executor import PostgreSQLExecutor
351355
352356 @pytest.fixture
353- def database (postgresql_proc ) :
357+ def database (postgresql_proc : PostgreSQLExecutor) -> Iterator[Connection] :
354358 # variable definition
355359
356360 janitor = DatabaseJanitor(
@@ -362,7 +366,7 @@ For this import DatabaseJanitor and use its init and drop methods:
362366 password = " secret_password" ,
363367 )
364368 janitor.init()
365- yield psycopg2 .connect(
369+ yield psycopg .connect(
366370 dbname = " my_test_database" ,
367371 user = postgresql_proc.user,
368372 password = " secret_password" ,
@@ -375,11 +379,14 @@ or use it as a context manager:
375379
376380.. code-block :: python
377381
382+ import psycopg
383+ from psycopg import Connection
378384 import pytest
379385 from pytest_postgresql.janitor import DatabaseJanitor
386+ from pytest_postgresql.executor import PostgreSQLExecutor
380387
381388 @pytest.fixture
382- def database (postgresql_proc ) :
389+ def database (postgresql_proc : PostgreSQLExecutor) -> Iterator[Connection] :
383390 # variable definition
384391
385392 with DatabaseJanitor(
@@ -390,7 +397,7 @@ or use it as a context manager:
390397 version = postgresql_proc.version,
391398 password = " secret_password" ,
392399 ):
393- yield psycopg2 .connect(
400+ yield psycopg .connect(
394401 dbname = " my_test_database" ,
395402 user = postgresql_proc.user,
396403 password = " secret_password" ,
@@ -425,14 +432,14 @@ In tests, make sure that all your tests are using **postgresql_noproc** fixture
425432
426433.. code-block :: python
427434
435+ from psycopg import Connection
428436 from pytest_postgresql import factories
429437
430-
431438 postgresql_in_docker = factories.postgresql_noproc()
432439 postgresql = factories.postgresql(" postgresql_in_docker" , dbname = " test" )
433440
434441
435- def test_postgres_docker (postgresql ) :
442+ def test_postgres_docker (postgresql : Connection) -> None :
436443 """ Run test."""
437444 cur = postgresql.cursor()
438445 cur.execute(" CREATE TABLE test (id serial PRIMARY KEY, num integer, data varchar);" )
@@ -453,9 +460,10 @@ your custom postgresql process fixture:
453460
454461.. code-block :: python
455462
463+ import psycopg
456464 import pytest_postgresql.factories
457- def load_database (** kwargs ) :
458- db_connection: connection = psycopg2 .connect(** kwargs)
465+ def load_database (** kwargs : str ) -> None :
466+ db_connection: connection = psycopg .connect(** kwargs)
459467 with db_connection.cursor() as cur:
460468 cur.execute(" CREATE TABLE stories (id serial PRIMARY KEY, name varchar);" )
461469 cur.execute(
@@ -486,7 +494,11 @@ How to use SQLAlchemy for common initialisation:
486494
487495.. code-block :: python
488496
489- def load_database (** kwargs ):
497+ import psycopg
498+ from sqlalchemy.orm import Session
499+
500+ def load_database (** kwargs ) -> None :
501+ from your_package import Base
490502 connection = f " postgresql+psycopg2:// { kwargs[' user' ]} :@ { kwargs[' host' ]} : { kwargs[' port' ]} / { kwargs[' dbname' ]} "
491503 engine = create_engine(connection)
492504 Base.metadata.create_all(engine)
@@ -499,7 +511,7 @@ How to use SQLAlchemy for common initialisation:
499511 postgresql = factories.postgresql(' postgresql_proc' ) # still need to check if this is actually needed or not
500512
501513 @pytest.fixture
502- def dbsession (postgresql ) :
514+ def dbsession (postgresql : psycopg.Connection) -> Session :
503515 connection = f ' postgresql+psycopg2:// { postgresql.info.user} :@ { postgresql.info.host} : { postgresql.info.port} / { postgresql.info.dbname} '
504516 engine = create_engine(connection)
505517
0 commit comments