[GitHub] [airflow] kaxil commented on a change in pull request #7516: [AIRFLOW-6894] Prevent DB query in example_dags

[GitBox <gi...@apache.org>]

kaxil commented on a change in pull request #7516: [AIRFLOW-6894] Prevent DB query in example_dags
URL: https://github.com/apache/airflow/pull/7516#discussion_r383042446
 
 

 ##########
 File path: tests/test_utils/asserts.py
 ##########
 @@ -16,9 +16,53 @@
 # under the License.
 
 import re
+from contextlib import contextmanager
+
+from sqlalchemy import event
+
+from airflow.settings import engine
 
 
 def assert_equal_ignore_multiple_spaces(case, first, second, msg=None):
     def _trim(s):
         return re.sub(r"\s+", " ", s.strip())
     return case.assertEqual(_trim(first), _trim(second), msg)
+
+
+class CountQueriesResult:
+    def __init__(self):
+        self.count = 0
+
+
+class CountQueries(object):
+    """
+    Counts the number of queries sent to Airflow Database in a given context.
+
+    Does not support multiple processes. When a new process is started in context, its queries will
+    not be included.
+    """
+    def __init__(self):
+        self.result = CountQueriesResult()
+
+    def __enter__(self):
+        event.listen(engine, "after_cursor_execute", self.after_cursor_execute)
+        return self.result
+
+    def __exit__(self, type, value, traceback):
+        event.remove(engine, "after_cursor_execute", self.after_cursor_execute)
+
+    def after_cursor_execute(self, *args, **kwargs):
+        self.result.count += 1
+
+
+count_queries = CountQueries
+
+
+@contextmanager
+def assert_queries_count(expected_count, message_fmt=None):
 
 Review comment:
   awesome

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services