diff --git a/edb/testbase/server.py b/edb/testbase/server.py index 89bf9d54393..408c89a3d6a 100644 --- a/edb/testbase/server.py +++ b/edb/testbase/server.py @@ -1045,7 +1045,6 @@ class DatabaseTestCase(ConnectedTestCase): async def setup_and_connect(cls): dbname = cls.get_database_name() - cls.admin_conn = None cls.con = None class_set_up = os.environ.get('EDGEDB_TEST_CASES_SET_UP', 'run') @@ -1053,16 +1052,17 @@ async def setup_and_connect(cls): # Only open an extra admin connection if necessary. if class_set_up == 'run': script = f'CREATE DATABASE {dbname};' - cls.admin_conn = await cls.connect() - await cls.admin_conn.execute(script) + admin_conn = await cls.connect() + await admin_conn.execute(script) + await admin_conn.aclose() elif class_set_up == 'inplace': dbname = edgedb_defines.EDGEDB_SUPERUSER_DB elif cls.uses_database_copies(): - cls.admin_conn = await cls.connect() + admin_conn = await cls.connect() - orig_testmode = await cls.admin_conn.query( + orig_testmode = await admin_conn.query( 'SELECT cfg::Config.__internal_testmode', ) if not orig_testmode: @@ -1072,7 +1072,7 @@ async def setup_and_connect(cls): # Enable testmode to unblock the template database syntax below. if not orig_testmode: - await cls.admin_conn.execute( + await admin_conn.execute( 'CONFIGURE SESSION SET __internal_testmode := true;', ) @@ -1087,7 +1087,7 @@ async def create_db(): timeout=30, ): async with tr: - await cls.admin_conn.execute( + await admin_conn.execute( f''' CREATE DATABASE {qlquote.quote_ident(dbname)} FROM {qlquote.quote_ident(base_db_name)} @@ -1096,10 +1096,12 @@ async def create_db(): await create_db() if not orig_testmode: - await cls.admin_conn.execute( + await admin_conn.execute( 'CONFIGURE SESSION SET __internal_testmode := false;', ) + await admin_conn.aclose() + cls.con = await cls.connect(database=dbname) if class_set_up != 'skip': @@ -1126,19 +1128,18 @@ async def teardown_and_disconnect(cls): if class_set_up == 'inplace': await cls.tearDownSingleDB() finally: - try: - await cls.con.aclose() + await cls.con.aclose() - if class_set_up == 'inplace': - pass - - elif class_set_up == 'run' or cls.uses_database_copies(): - dbname = qlquote.quote_ident(cls.get_database_name()) - await drop_db(cls.admin_conn, dbname) + if class_set_up == 'inplace': + pass - finally: - if cls.admin_conn is not None: - await cls.admin_conn.aclose() + elif class_set_up == 'run' or cls.uses_database_copies(): + dbname = qlquote.quote_ident(cls.get_database_name()) + admin_conn = await cls.connect() + try: + await drop_db(admin_conn, dbname) + finally: + await admin_conn.aclose() @classmethod def get_database_name(cls):