Differential D2184 Diff 7492 swh/storage/tests/conftest.py

Changeset View

Standalone View

swh/storage/tests/conftest.py

# Copyright (C) 2019 The Software Heritage developers		# Copyright (C) 2019 The Software Heritage developers
# See the AUTHORS file at the top-level directory of this distribution		# See the AUTHORS file at the top-level directory of this distribution
# License: GNU General Public License version 3, or any later version		# License: GNU General Public License version 3, or any later version
# See top-level LICENSE file for more information		# See top-level LICENSE file for more information

import glob		import glob
import pytest		import pytest

		from typing import Union

from pytest_postgresql import factories		from pytest_postgresql import factories
from pytest_postgresql.janitor import DatabaseJanitor, psycopg2		from pytest_postgresql.janitor import DatabaseJanitor, psycopg2, Version

from os import path, environ		from os import path, environ
from hypothesis import settings		from hypothesis import settings
from typing import Dict		from typing import Dict

import swh.storage		import swh.storage

from swh.core.utils import numfile_sortkey as sortkey		from swh.core.utils import numfile_sortkey as sortkey
▲ Show 20 Lines • Show All 48 Lines • ▼ Show 20 Lines	def swh_origins(swh_storage):
origins = gen_origins(n=100)		origins = gen_origins(n=100)
swh_storage.origin_add(origins)		swh_storage.origin_add(origins)
return origins		return origins


# the postgres_fact factory fixture below is mostly a copy of the code		# 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		# from pytest-postgresql. We need a custom version here to be able to
# specify our version of the DBJanitor we use.		# specify our version of the DBJanitor we use.
def postgresql_fact(process_fixture_name, db_name=None):		def postgresql_fact(process_fixture_name, db_name=None, dump_files=DUMP_FILES):
@pytest.fixture		@pytest.fixture
def postgresql_factory(request):		def postgresql_factory(request):
"""		"""
Fixture factory for PostgreSQL.		Fixture factory for PostgreSQL.

:param FixtureRequest request: fixture request object		:param FixtureRequest request: fixture request object
:rtype: psycopg2.connection		:rtype: psycopg2.connection
:returns: postgresql client		:returns: postgresql client
"""		"""
config = factories.get_config(request)		config = factories.get_config(request)
if not psycopg2:		if not psycopg2:
raise ImportError(		raise ImportError(
'No module named psycopg2. Please install it.'		'No module named psycopg2. Please install it.'
)		)
proc_fixture = request.getfixturevalue(process_fixture_name)		proc_fixture = request.getfixturevalue(process_fixture_name)

# _, config = try_import('psycopg2', request)		# _, config = try_import('psycopg2', request)
pg_host = proc_fixture.host		pg_host = proc_fixture.host
pg_port = proc_fixture.port		pg_port = proc_fixture.port
pg_user = proc_fixture.user		pg_user = proc_fixture.user
pg_options = proc_fixture.options		pg_options = proc_fixture.options
pg_db = db_name or config['dbname']		pg_db = db_name or config['dbname']

with SwhDatabaseJanitor(		with SwhDatabaseJanitor(
pg_user, pg_host, pg_port, pg_db, proc_fixture.version		pg_user, pg_host, pg_port, pg_db, proc_fixture.version,
		dump_files=dump_files
):		):
connection = psycopg2.connect(		connection = psycopg2.connect(
dbname=pg_db,		dbname=pg_db,
user=pg_user,		user=pg_user,
host=pg_host,		host=pg_host,
port=pg_port,		port=pg_port,
options=pg_options		options=pg_options
)		)
yield connection		yield connection
connection.close()		connection.close()

return postgresql_factory		return postgresql_factory


swh_storage_postgresql = postgresql_fact('postgresql_proc')		swh_storage_postgresql = postgresql_fact('postgresql_proc')


# This version of the DatabaseJanitor implement a different setup/teardown		# This version of the DatabaseJanitor implement a different setup/teardown
# behavior than than the stock one: instead of droping, creating and		# behavior than than the stock one: instead of droping, creating and
# initializing the database for each test, it create and initialize the db only		# initializing the database for each test, it create and initialize the db only
# once, then it truncate the tables. This is needed to have acceptable test		# once, then it truncate the tables. This is needed to have acceptable test
# performances.		# performances.
class SwhDatabaseJanitor(DatabaseJanitor):		class SwhDatabaseJanitor(DatabaseJanitor):
		def __init__(
		self,
		user: str,
		host: str,
		port: str,
		db_name: str,
		version: Union[str, float, Version],
		dump_files: str = DUMP_FILES
		) -> None:
		super().__init__(user, host, port, db_name, version)
		self.dump_files = sorted(
		glob.glob(dump_files), key=sortkey)

def db_setup(self):		def db_setup(self):
with psycopg2.connect(		with psycopg2.connect(
dbname=self.db_name,		dbname=self.db_name,
user=self.user,		user=self.user,
host=self.host,		host=self.host,
port=self.port,		port=self.port,
) as cnx:		) as cnx:
with cnx.cursor() as cur:		with cnx.cursor() as cur:
all_dump_files = sorted(		for fname in self.dump_files:
glob.glob(DUMP_FILES), key=sortkey)
for fname in all_dump_files:
with open(fname) as fobj:		with open(fname) as fobj:
sql = fobj.read().replace('concurrently', '')		sql = fobj.read().replace('concurrently', '').strip()
		if sql:
cur.execute(sql)		cur.execute(sql)
cnx.commit()		cnx.commit()

def db_reset(self):		def db_reset(self):
with psycopg2.connect(		with psycopg2.connect(
dbname=self.db_name,		dbname=self.db_name,
user=self.user,		user=self.user,
host=self.host,		host=self.host,
port=self.port,		port=self.port,
▲ Show 20 Lines • Show All 108 Lines • Show Last 20 Lines