diff options
Diffstat (limited to 'lib/sqlalchemy/testing/assertions.py')
| -rw-r--r-- | lib/sqlalchemy/testing/assertions.py | 23 |
1 files changed, 12 insertions, 11 deletions
diff --git a/lib/sqlalchemy/testing/assertions.py b/lib/sqlalchemy/testing/assertions.py index 5dd32e8cc..c66ba71c3 100644 --- a/lib/sqlalchemy/testing/assertions.py +++ b/lib/sqlalchemy/testing/assertions.py @@ -906,7 +906,17 @@ class AssertsExecutionResults: db, callable_, assertsql.CountStatements(count) ) - def assert_multiple_sql_count(self, dbs, callable_, counts): + @contextlib.contextmanager + def assert_execution(self, db, *rules): + with self.sql_execution_asserter(db) as asserter: + yield + asserter.assert_(*rules) + + def assert_statement_count(self, db, count): + return self.assert_execution(db, assertsql.CountStatements(count)) + + @contextlib.contextmanager + def assert_statement_count_multi_db(self, dbs, counts): recs = [ (self.sql_execution_asserter(db), db, count) for (db, count) in zip(dbs, counts) @@ -915,21 +925,12 @@ class AssertsExecutionResults: for ctx, db, count in recs: asserters.append(ctx.__enter__()) try: - return callable_() + yield finally: for asserter, (ctx, db, count) in zip(asserters, recs): ctx.__exit__(None, None, None) asserter.assert_(assertsql.CountStatements(count)) - @contextlib.contextmanager - def assert_execution(self, db, *rules): - with self.sql_execution_asserter(db) as asserter: - yield - asserter.assert_(*rules) - - def assert_statement_count(self, db, count): - return self.assert_execution(db, assertsql.CountStatements(count)) - class ComparesIndexes: def compare_table_index_with_expected( |