kaxil commented on a change in pull request #9509:
URL: https://github.com/apache/airflow/pull/9509#discussion_r445487074
##########
File path: tests/models/test_taskinstance.py
##########
@@ -1699,3 +1704,54 @@ def test_refresh_from_task(pool_override):
assert ti.max_tries == task.retries
assert ti.executor_config == task.executor_config
assert ti.operator == DummyOperator.__name__
+
+
+class TestRunRawTaskQueriesCount(unittest.TestCase):
+ """
+ These tests are designed to detect changes in the number of queries for
+ different DAG files. These tests allow easy detection when a change is
+ made that affects the performance of the SchedulerJob.
+ """
+
+ @staticmethod
+ def _clean():
+ clear_db_runs()
+ clear_db_pools()
+ clear_db_dags()
+ clear_db_sla_miss()
+ clear_db_errors()
+
+ def setUp(self) -> None:
+ self._clean()
+
+ def tearDown(self) -> None:
+ self._clean()
+
+ @parameterized.expand([
+ # Expected queries, mark_success
+ (7, False),
+ (5, True),
+ ])
+ def test_execute_queries_count(self, expected_query_count, mark_success):
+ with create_session() as s:
+ dag = DAG('test_queries', start_date=DEFAULT_DATE)
+ task = DummyOperator(task_id='op', dag=dag)
+ ti = TI(task=task, execution_date=datetime.datetime.now())
+ ti.state = State.RUNNING
+ s.merge(ti)
+
+ with assert_queries_count(expected_query_count):
+ ti._run_raw_task(mark_success=mark_success)
+
+ def test_execute_queries_count_store_serialized(self):
+ with create_session() as s:
+ dag = DAG('test_queries', start_date=DEFAULT_DATE)
+ task = DummyOperator(task_id='op', dag=dag)
+ ti = TI(task=task, execution_date=datetime.datetime.now())
+ ti.state = State.RUNNING
+ s.merge(ti)
+
+ with assert_queries_count(10), mock.patch(
Review comment:
```suggestion
with assert_queries_count(10), patch(
```
##########
File path: tests/models/test_taskinstance.py
##########
@@ -22,6 +22,7 @@
import unittest
import urllib
from typing import List, Optional, Union, cast
+from unittest import mock
Review comment:
```suggestion
```
----------------------------------------------------------------
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:
[email protected]
