- add a nifty contextmanager form of assert_execution/assert_statement_count

1 files changed, 11 insertions, 0 deletions

diff --git a/lib/sqlalchemy/testing/assertions.py b/lib/sqlalchemy/testing/assertions.py
index 2b305b98a..e74d13a97 100644
--- a/lib/sqlalchemy/testing/assertions.py
+++ b/lib/sqlalchemy/testing/assertions.py
@@ -14,6 +14,7 @@ from . import assertsql
 from . import config
 import itertools
 from .util import fail
+import contextlib
 
 def emits_warning(*messages):
     """Mark a test as emitting a warning.
@@ -348,4 +349,14 @@ class AssertsExecutionResults(object):
     def assert_sql_count(self, db, callable_, count):
         self.assert_sql_execution(db, callable_, assertsql.CountStatements(count))
 
+    @contextlib.contextmanager
+    def assert_execution(self, *rules):
+        assertsql.asserter.add_rules(rules)
+        try:
+            yield
+            assertsql.asserter.statement_complete()
+        finally:
+            assertsql.asserter.clear_rules()
 
+    def assert_statement_count(self, count):
+        return self.assert_execution(assertsql.CountStatements(count))
\ No newline at end of file