diff --git a/swh/storage/tests/test_storage.py b/swh/storage/tests/test_storage.py --- a/swh/storage/tests/test_storage.py +++ b/swh/storage/tests/test_storage.py @@ -34,11 +34,6 @@ class StorageTestDbFixture(StorageTestFixture): def setUp(self): super().setUp() - - db = self.test_db[self.TEST_DB_NAME] - self.conn = db.conn - self.cursor = db.cursor - self.maxDiff = None def tearDown(self): @@ -711,13 +706,15 @@ if hasattr(self.storage, 'objstorage'): self.assertIn(cont['sha1'], self.storage.objstorage) - self.cursor.execute('SELECT sha1, sha1_git, sha256, length, status' - ' FROM content WHERE sha1 = %s', - (cont['sha1'],)) - datum = self.cursor.fetchone() + + with self.db_transaction() as (_, cur): + cur.execute('SELECT sha1, sha1_git, sha256, length, status' + ' FROM content WHERE sha1 = %s', + (cont['sha1'],)) + datum = cur.fetchone() + self.assertEqual( - (datum[0].tobytes(), datum[1].tobytes(), datum[2].tobytes(), - datum[3], datum[4]), + datum, (cont['sha1'], cont['sha1_git'], cont['sha256'], cont['length'], 'visible')) @@ -802,13 +799,15 @@ if hasattr(self.storage, 'objstorage'): self.assertNotIn(cont['sha1'], self.storage.objstorage) - self.cursor.execute('SELECT sha1, sha1_git, sha256, length, status' - ' FROM content WHERE sha1 = %s', - (cont['sha1'],)) - datum = self.cursor.fetchone() + + with self.db_transaction() as (_, cur): + cur.execute('SELECT sha1, sha1_git, sha256, length, status' + ' FROM content WHERE sha1 = %s', + (cont['sha1'],)) + datum = cur.fetchone() + self.assertEqual( - (datum[0].tobytes(), datum[1].tobytes(), datum[2].tobytes(), - datum[3], datum[4]), + datum, (cont['sha1'], cont['sha1_git'], cont['sha256'], cont['length'], 'visible')) @@ -844,26 +843,23 @@ 'skipped_content:add': 2, }) - self.cursor.execute('SELECT sha1, sha1_git, sha256, blake2s256, ' - 'length, status, reason ' - 'FROM skipped_content ORDER BY sha1_git') + with self.db_transaction() as (_, cur): + cur.execute('SELECT sha1, sha1_git, sha256, blake2s256, ' + 'length, status, reason ' + 'FROM skipped_content ORDER BY sha1_git') - datums = self.cursor.fetchall() + data = cur.fetchall() - self.assertEqual(2, len(datums)) - datum = datums[0] + self.assertEqual(2, len(data)) self.assertEqual( - (datum[0].tobytes(), datum[1].tobytes(), datum[2].tobytes(), - datum[3].tobytes(), datum[4], datum[5], datum[6]), + data[0], (cont['sha1'], cont['sha1_git'], cont['sha256'], cont['blake2s256'], cont['length'], 'absent', 'Content too long') ) - datum2 = datums[1] self.assertEqual( - (datum2[0].tobytes(), datum2[1].tobytes(), datum2[2].tobytes(), - datum2[3], datum2[4], datum2[5], datum2[6]), + data[1], (cont2['sha1'], cont2['sha1_git'], cont2['sha256'], cont2['blake2s256'], cont2['length'], 'absent', 'Content too long') @@ -4056,18 +4052,13 @@ results = queue.Queue() def thread(): - db = None try: - db = self.storage.get_db() - with db.transaction() as cur: + with self.db_transaction() as (db, cur): ret = self.storage.content_add([self.cont], db=db, cur=cur) results.put((threading.get_ident(), 'data', ret)) except Exception as e: results.put((threading.get_ident(), 'exc', e)) - finally: - if db: - self.storage.put_db(db) t1 = threading.Thread(target=thread) t2 = threading.Thread(target=thread) @@ -4116,7 +4107,7 @@ self.storage.content_update([cont], keys=['sha1_git']) - with self.storage.get_db().transaction() as cur: + with self.db_transaction() as (_, cur): cur.execute('SELECT sha1, sha1_git, sha256, length, status' ' FROM content WHERE sha1 = %s', (cont['sha1'],)) @@ -4131,7 +4122,7 @@ def test_content_update_with_new_cols(self): self.storage.journal_writer = None # TODO, not supported - with self.storage.get_db().transaction() as cur: + with self.db_transaction() as (db, cur): cur.execute("""alter table content add column test text default null, add column test2 text default null""") @@ -4142,10 +4133,11 @@ cont['test2'] = 'value-2' self.storage.content_update([cont], keys=['test', 'test2']) - with self.storage.get_db().transaction() as cur: + with self.db_transaction() as (_, cur): cur.execute( - 'SELECT sha1, sha1_git, sha256, length, status, test, test2' - ' FROM content WHERE sha1 = %s', + '''SELECT sha1, sha1_git, sha256, length, status, + test, test2 + FROM content WHERE sha1 = %s''', (cont['sha1'],)) datum = cur.fetchone() @@ -4156,6 +4148,6 @@ (cont['sha1'], cont['sha1_git'], cont['sha256'], cont['length'], 'visible', cont['test'], cont['test2'])) - with self.storage.get_db().transaction() as cur: + with self.db_transaction() as (_, cur): cur.execute("""alter table content drop column test, - drop column test2""") + drop column test2""")