diff options
Diffstat (limited to 'lib/sqlalchemy/testing/assertions.py')
| -rw-r--r-- | lib/sqlalchemy/testing/assertions.py | 13 |
1 files changed, 5 insertions, 8 deletions
diff --git a/lib/sqlalchemy/testing/assertions.py b/lib/sqlalchemy/testing/assertions.py index c0854ea55..08d0f0aac 100644 --- a/lib/sqlalchemy/testing/assertions.py +++ b/lib/sqlalchemy/testing/assertions.py @@ -520,13 +520,10 @@ class AssertsExecutionResults(object): db, callable_, assertsql.CountStatements(count)) @contextlib.contextmanager - def assert_execution(self, *rules): - assertsql.asserter.add_rules(rules) - try: + def assert_execution(self, db, *rules): + with self.sql_execution_asserter(db) as asserter: yield - assertsql.asserter.statement_complete() - finally: - assertsql.asserter.clear_rules() + asserter.assert_(*rules) - def assert_statement_count(self, count): - return self.assert_execution(assertsql.CountStatements(count)) + def assert_statement_count(self, db, count): + return self.assert_execution(db, assertsql.CountStatements(count)) |