-
Notifications
You must be signed in to change notification settings - Fork 50
Expand file tree
/
Copy pathconftest.py
More file actions
67 lines (53 loc) · 2.38 KB
/
conftest.py
File metadata and controls
67 lines (53 loc) · 2.38 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
import asyncio
import os
import random
import pytest
import sqlalchemy
import sqlalchemy.ext.asyncio
@pytest.fixture(scope="session")
def postgres_uri() -> str:
pg_host = os.environ.get("PG_HOST", "postgres")
pg_port = os.environ.get("PG_PORT", 5432)
pg_user = os.environ.get("PG_USER", "postgres")
pg_password = os.environ.get("PG_PASSWORD", "mysecretpassword")
pg_db = os.environ.get("PG_DATABASE", "dinotest")
return f"postgresql://{pg_user}:{pg_password}@{pg_host}:{pg_port}/{pg_db}"
@pytest.fixture(scope="session")
def sqlalchemy_connection(postgres_uri) -> sqlalchemy.engine.Connection:
engine = sqlalchemy.create_engine(postgres_uri, future=True)
with engine.connect() as conn:
yield conn
@pytest.fixture(scope="function")
def db(sqlalchemy_connection: sqlalchemy.engine.Connection) -> sqlalchemy.engine.Connection:
conn = sqlalchemy_connection
schema_name = f"sqltest_{random.randint(0, 1000)}"
conn.execute(sqlalchemy.text(f"CREATE SCHEMA {schema_name}"))
conn.execute(sqlalchemy.text(f"SET search_path TO {schema_name}"))
conn.commit()
yield conn
conn.rollback()
conn.execute(sqlalchemy.text(f"DROP SCHEMA {schema_name} CASCADE"))
conn.execute(sqlalchemy.text("SET search_path TO public"))
@pytest.fixture(scope="session")
async def async_sqlalchemy_connection(postgres_uri) -> sqlalchemy.ext.asyncio.AsyncConnection:
postgres_uri = postgres_uri.replace("postgresql", "postgresql+asyncpg")
engine = sqlalchemy.ext.asyncio.create_async_engine(postgres_uri)
async with engine.connect() as conn:
yield conn
@pytest.fixture(scope="function")
async def async_db(async_sqlalchemy_connection: sqlalchemy.ext.asyncio.AsyncConnection) -> sqlalchemy.ext.asyncio.AsyncConnection:
conn = async_sqlalchemy_connection
schema_name = f"sqltest_{random.randint(0, 1000)}"
await conn.execute(sqlalchemy.text(f"CREATE SCHEMA {schema_name}"))
await conn.execute(sqlalchemy.text(f"SET search_path TO {schema_name}"))
await conn.commit()
yield conn
await conn.rollback()
await conn.execute(sqlalchemy.text(f"DROP SCHEMA {schema_name} CASCADE"))
await conn.execute(sqlalchemy.text("SET search_path TO public"))
@pytest.fixture(scope="session")
def event_loop():
"""Change event_loop fixture to session level."""
loop = asyncio.get_event_loop_policy().new_event_loop()
yield loop
loop.close()