-
-
Notifications
You must be signed in to change notification settings - Fork 21
/
conftest.py
130 lines (96 loc) · 3.79 KB
/
conftest.py
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
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
import os
import random
import string
import time
from contextlib import contextmanager
from typing import Generator
import pytest
from cryptography.fernet import Fernet
from flask import Flask
from flask.testing import FlaskClient
from pytest_mock import MockFixture
from sqlalchemy import create_engine, text
from sqlalchemy.exc import OperationalError
from sqlalchemy.orm import Session, sessionmaker
from hushline import create_app, db
CONN_FMT_STR = "postgresql+psycopg:https://hushline:[email protected]:5432/{database}"
TEMPLATE_DB_NAME = "app_db_template"
def random_name(size: int) -> str:
return "".join([random.choice(string.ascii_lowercase) for _ in range(size)]) # noqa: S311
@contextmanager
def temp_session(conn_str: str) -> Generator[Session, None, None]:
engine = create_engine(conn_str)
engine = engine.execution_options(isolation_level="AUTOCOMMIT")
session = sessionmaker(bind=engine)()
yield session
# aggressively terminate all connections
session.close()
session.connection().connection.invalidate()
engine.dispose()
@pytest.fixture(scope="session")
def _db_template() -> Generator[None, None, None]:
"""A template database that all other databases are created from.
Effectively allows caching of `db.create_all()`"""
conn_str = CONN_FMT_STR.format(database="hushline")
counter = 5
while True:
try:
with temp_session(conn_str) as session:
session.execute(text(f"CREATE DATABASE {TEMPLATE_DB_NAME} WITH TEMPLATE template1"))
break
except OperationalError:
if counter > 0:
counter -= 1
time.sleep(1)
else:
raise
db_uri = CONN_FMT_STR.format(database=TEMPLATE_DB_NAME)
init_db_via_create_all(db_uri)
yield
with temp_session(conn_str) as session:
session.execute(text(f"DROP DATABASE {TEMPLATE_DB_NAME}"))
def init_db_via_create_all(db_uri: str) -> None:
# dumb hack to easily get the create_all() functionality
os.environ["SQLALCHEMY_DATABASE_URI"] = db_uri
app = create_app()
with app.app_context():
db.session.commit()
db.create_all()
db.session.close()
db.session.connection().connection.invalidate() # type: ignore
@pytest.fixture()
def database(_db_template: None) -> str:
"""A clean Postgres database from the template with DDLs applied"""
db_name = random_name(16)
conn_str = CONN_FMT_STR.format(database="hushline")
engine = create_engine(conn_str)
engine = engine.execution_options(isolation_level="AUTOCOMMIT")
session = sessionmaker(bind=engine)()
sql = text(f"CREATE DATABASE {db_name} WITH TEMPLATE {TEMPLATE_DB_NAME}")
session.execute(sql)
# aggressively terminate all connections
session.close()
session.connection().connection.invalidate()
engine.dispose()
print(f"Postgres DB: {db_name}, template: {TEMPLATE_DB_NAME}") # to help with debugging tests
return db_name
@pytest.fixture()
def _config(mocker: MockFixture) -> None:
mocker.patch.dict(os.environ, {})
@pytest.fixture()
def app(_config: None, database: str) -> Generator[Flask, None, None]:
os.environ["SQLALCHEMY_TRACK_MODIFICATIONS"] = "False"
os.environ["REGISTRATION_CODES_REQUIRED"] = "False"
os.environ["ENCRYPTION_KEY"] = Fernet.generate_key().decode()
os.environ["SQLALCHEMY_DATABASE_URI"] = CONN_FMT_STR.format(database=database)
app = create_app()
app.config["TESTING"] = True
app.config["WTF_CSRF_ENABLED"] = False
app.config["SERVER_NAME"] = "localhost:8080"
app.config["PREFERRED_URL_SCHEME"] = "http"
with app.app_context():
yield app
@pytest.fixture()
def client(app: Flask) -> Generator[FlaskClient, None, None]:
with app.test_client() as client:
yield client