Page MenuHomeSoftware Heritage
Files F9696502

pytest_plugin.py
No OneTemporary
Actions

pytest_plugin.py
View Options

# Copyright (C) 2020 The Software Heritage developers
# See the AUTHORS file at the top-level directory of this distribution
# License: GNU General Public License version 3, or any later version
# See top-level LICENSE file for more information
import glob
from importlib import import_module
import logging
import subprocess
from typing import List, Optional, Set, Union
from _pytest.fixtures import FixtureRequest
import psycopg2
import pytest
from pytest_postgresql.janitor import DatabaseJanitor
from swh.core.utils import numfile_sortkey as sortkey
# to keep mypy happy regardless pytest-postgresql version
try:
_pytest_pgsql_get_config_module = import_module("pytest_postgresql.config")
except ImportError:
# pytest_postgresql < 3.0.0
_pytest_pgsql_get_config_module = import_module("pytest_postgresql.factories")
_pytest_postgresql_get_config = getattr(_pytest_pgsql_get_config_module, "get_config")
logger = logging.getLogger(__name__)
class SWHDatabaseJanitor(DatabaseJanitor):
"""SWH database janitor implementation with a a different setup/teardown policy than
than the stock one. Instead of dropping, creating and initializing the database for
each test, it creates and initializes the db once, then truncates the tables (and
sequences) in between tests.
This is needed to have acceptable test performances.
"""
def __init__(
self,
user: str,
host: str,
port: str,
dbname: str,
version: Union[str, float],
dump_files: Union[None, str, List[str]] = None,
no_truncate_tables: Set[str] = set(),
) -> None:
super().__init__(user, host, port, dbname, version)
if not hasattr(self, "dbname") and hasattr(self, "db_name"):
# pytest_postgresql < 3.0.0
self.dbname = getattr(self, "db_name")
if dump_files is None:
self.dump_files = []
elif isinstance(dump_files, str):
self.dump_files = sorted(glob.glob(dump_files), key=sortkey)
else:
self.dump_files = dump_files
# do no truncate the following tables
self.no_truncate_tables = set(no_truncate_tables)
def db_setup(self):
conninfo = (
f"host={self.host} user={self.user} port={self.port} dbname={self.dbname}"
)
for fname in self.dump_files:
subprocess.check_call(
[
"psql",
"--quiet",
"--no-psqlrc",
"-v",
"ON_ERROR_STOP=1",
"-d",
conninfo,
"-f",
fname,
]
)
def db_reset(self):
"""Truncate tables (all but self.no_truncate_tables set) and sequences
"""
with psycopg2.connect(
dbname=self.dbname, user=self.user, host=self.host, port=self.port,
) as cnx:
with cnx.cursor() as cur:
cur.execute(
"SELECT table_name FROM information_schema.tables "
"WHERE table_schema = %s",
("public",),
)
all_tables = set(table for (table,) in cur.fetchall())
tables_to_truncate = all_tables - self.no_truncate_tables
for table in tables_to_truncate:
cur.execute("TRUNCATE TABLE %s CASCADE" % table)
cur.execute(
"SELECT sequence_name FROM information_schema.sequences "
"WHERE sequence_schema = %s",
("public",),
)
seqs = set(seq for (seq,) in cur.fetchall())
for seq in seqs:
cur.execute("ALTER SEQUENCE %s RESTART;" % seq)
cnx.commit()
def init(self):
"""Initialize db. Create the db if it does not exist. Reset it if it exists."""
with self.cursor() as cur:
cur.execute(
"SELECT COUNT(1) FROM pg_database WHERE datname=%s;", (self.dbname,)
)
db_exists = cur.fetchone()[0] == 1
if db_exists:
cur.execute(
"UPDATE pg_database SET datallowconn=true WHERE datname = %s;",
(self.dbname,),
)
self.db_reset()
return
# initialize the inexistent db
with self.cursor() as cur:
cur.execute('CREATE DATABASE "{}";'.format(self.dbname))
self.db_setup()
def drop(self):
"""The original DatabaseJanitor implementation prevents new connections from happening,
destroys current opened connections and finally drops the database.
We actually do not want to drop the db so we instead do nothing and resets
(truncate most tables and sequences) the db instead, in order to have some
acceptable performance.
"""
pass
# the postgres_fact factory fixture below is mostly a copy of the code
# from pytest-postgresql. We need a custom version here to be able to
# specify our version of the DBJanitor we use.
def postgresql_fact(
process_fixture_name: str,
dbname: Optional[str] = None,
dump_files: Union[str, List[str]] = "",
no_truncate_tables: Set[str] = {"dbversion"},
):
@pytest.fixture
def postgresql_factory(request: FixtureRequest):
"""Fixture factory for PostgreSQL.
:param FixtureRequest request: fixture request object
:rtype: psycopg2.connection
:returns: postgresql client
"""
config = _pytest_postgresql_get_config(request)
proc_fixture = request.getfixturevalue(process_fixture_name)
pg_host = proc_fixture.host
pg_port = proc_fixture.port
pg_user = proc_fixture.user
pg_options = proc_fixture.options
pg_db = dbname or config["dbname"]
with SWHDatabaseJanitor(
pg_user,
pg_host,
pg_port,
pg_db,
proc_fixture.version,
dump_files=dump_files,
no_truncate_tables=no_truncate_tables,
):
connection = psycopg2.connect(
dbname=pg_db,
user=pg_user,
host=pg_host,
port=pg_port,
options=pg_options,
)
yield connection
connection.close()
return postgresql_factory

File Metadata

Mime Type
text/x-python
Expires
Mon, Aug 18, 8:18 PM (5 d, 17 h ago)
Storage Engine
blob
Storage Format
Raw Data
Storage Handle
3278345

Event Timeline

About · Developer information · Legal