From 2c6a7dfcbfc6e3782c5d2d88b28a54efca8aff8b Mon Sep 17 00:00:00 2001 From: Patrick Cloke Date: Thu, 9 Nov 2023 16:19:42 -0500 Subject: Use attempt_to_set_autocommit everywhere. (#16615) To avoid asserting the type of the database connection. --- tests/server.py | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) (limited to 'tests/server.py') diff --git a/tests/server.py b/tests/server.py index 5a63ecee9f..2b63ed3dd8 100644 --- a/tests/server.py +++ b/tests/server.py @@ -88,7 +88,7 @@ from synapse.module_api.callbacks.third_party_event_rules_callbacks import ( from synapse.server import HomeServer from synapse.storage import DataStore from synapse.storage.database import LoggingDatabaseConnection -from synapse.storage.engines import PostgresEngine, create_engine +from synapse.storage.engines import create_engine from synapse.storage.prepare_database import prepare_database from synapse.types import ISynapseReactor, JsonDict from synapse.util import Clock @@ -1029,9 +1029,7 @@ def setup_test_homeserver( # Create the database before we actually try and connect to it, based off # the template database we generate in setupdb() - if isinstance(db_engine, PostgresEngine): - import psycopg2.extensions - + if USE_POSTGRES_FOR_TESTS: db_conn = db_engine.module.connect( dbname=POSTGRES_BASE_DB, user=POSTGRES_USER, @@ -1039,8 +1037,7 @@ def setup_test_homeserver( port=POSTGRES_PORT, password=POSTGRES_PASSWORD, ) - assert isinstance(db_conn, psycopg2.extensions.connection) - db_conn.autocommit = True + db_engine.attempt_to_set_autocommit(db_conn, True) cur = db_conn.cursor() cur.execute("DROP DATABASE IF EXISTS %s;" % (test_db,)) cur.execute( @@ -1065,13 +1062,12 @@ def setup_test_homeserver( hs.setup() - if isinstance(db_engine, PostgresEngine): + if USE_POSTGRES_FOR_TESTS: database_pool = hs.get_datastores().databases[0] # We need to do cleanup on PostgreSQL def cleanup() -> None: import psycopg2 - import psycopg2.extensions # Close all the db pools database_pool._db_pool.close() @@ -1086,8 +1082,7 @@ def setup_test_homeserver( port=POSTGRES_PORT, password=POSTGRES_PASSWORD, ) - assert isinstance(db_conn, psycopg2.extensions.connection) - db_conn.autocommit = True + db_engine.attempt_to_set_autocommit(db_conn, True) cur = db_conn.cursor() # Try a few times to drop the DB. Some things may hold on to the -- cgit 1.4.1