From 49e3d183e5bca2c6e475649eda4b6d1a42323dae Mon Sep 17 00:00:00 2001
From: "maciej.dembiczak" <maciej.dembiczak@openx.com>
Date: Fri, 19 May 2023 09:49:54 +0200
Subject: [PATCH 1/3] assign filter output to variable

---
 airflow/models/taskinstance.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py
index bb7175a4a330e..2e416ce2eccf5 100644
--- a/airflow/models/taskinstance.py
+++ b/airflow/models/taskinstance.py
@@ -2551,7 +2551,9 @@ def get_num_running_task_instances(self, session: Session, same_dagrun=False) ->
             TaskInstance.state == State.RUNNING,
         )
         if same_dagrun:
-            num_running_task_instances_query.filter(TaskInstance.run_id == self.run_id)
+            num_running_task_instances_query = num_running_task_instances_query.filter(
+                TaskInstance.run_id == self.run_id
+            )
         return num_running_task_instances_query.scalar()
 
     def init_run_context(self, raw: bool = False) -> None:

From 2b1e74f6acb91a9337c188cbd1f6aa9df66bec11 Mon Sep 17 00:00:00 2001
From: "maciej.dembiczak" <maciej.dembiczak@openx.com>
Date: Tue, 30 May 2023 14:28:29 +0200
Subject: [PATCH 2/3] add tests for max_active_tis_per_dagrun

---
 tests/models/test_taskinstance.py | 86 +++++++++++++++++++++----------
 1 file changed, 58 insertions(+), 28 deletions(-)

diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
index 0dcfb632142de..9f62fa9833ca3 100644
--- a/tests/models/test_taskinstance.py
+++ b/tests/models/test_taskinstance.py
@@ -82,7 +82,7 @@
 from airflow.utils.db import merge_conn
 from airflow.utils.module_loading import qualname
 from airflow.utils.session import create_session, provide_session
-from airflow.utils.state import State, TaskInstanceState
+from airflow.utils.state import State, TaskInstanceState, DagRunState
 from airflow.utils.task_group import TaskGroup
 from airflow.utils.types import DagRunType
 from airflow.utils.xcom import XCOM_RETURN_KEY
@@ -1770,6 +1770,59 @@ def test_get_num_running_task_instances(self, create_task_instance):
         assert 1 == ti2.get_num_running_task_instances(session=session)
         assert 1 == ti3.get_num_running_task_instances(session=session)
 
+    def test_get_num_running_task_instances_per_dagrun(self, create_task_instance, dag_maker):
+        session = settings.Session()
+
+        @task
+        def dummy_mapped_task(x: int):
+            return x
+
+        with dag_maker(dag_id="test_dag") as dag:
+            dummy_mapped_task.expand(x=[1, 2])
+            dagrun1 = dag.create_dagrun(execution_date=timezone.utcnow(),
+                                       state=DagRunState.RUNNING,
+                                       run_id="run_id_1",
+                                       session=session)
+
+            ti1 = dagrun1.task_instances[0]
+            ti2 = dagrun1.task_instances[1]
+
+        with dag_maker(dag_id="test_dag") as dag:
+            dummy_mapped_task.expand(x=[3])
+            dagrun2 = dag.create_dagrun(execution_date=timezone.utcnow(),
+                                       state=DagRunState.RUNNING,
+                                       run_id="run_id_2",
+                                       session=session)
+
+            ti3 = dagrun2.task_instances[0]
+
+        with dag_maker(dag_id="test_dag") as dag:
+            dummy_mapped_task.expand(x=[4])
+            dagrun3 = dag.create_dagrun(execution_date=timezone.utcnow(),
+                                       state=DagRunState.RUNNING,
+                                       run_id="run_id_3",
+                                       session=session)
+
+            ti4 = dagrun3.task_instances[0]
+
+        assert ti1 in session
+        assert ti2 in session
+        assert ti3 in session
+        assert ti4 in session
+
+        ti1.state = State.RUNNING
+        ti2.state = State.QUEUED
+        ti3.state = State.RUNNING
+        ti4.state = State.RUNNING
+
+        session.commit()
+
+        assert 1 == ti1.get_num_running_task_instances(session=session, same_dagrun=True)
+        assert 1 == ti2.get_num_running_task_instances(session=session, same_dagrun=True)
+        assert 3 == ti3.get_num_running_task_instances(session=session)
+        assert 1 == ti4.get_num_running_task_instances(session=session, same_dagrun=True)
+
+
     def test_log_url(self, create_task_instance):
         ti = create_task_instance(dag_id="dag", task_id="op", execution_date=timezone.datetime(2018, 1, 1))
 
@@ -1903,7 +1956,6 @@ def test_email_alert_with_filenotfound_config(self, mock_send_email, dag_maker):
     @patch("airflow.models.taskinstance.send_email")
     def test_failure_mapped_taskflow(self, mock_send_email, dag_maker, session, task_id):
         with dag_maker(session=session) as dag:
-
             @dag.task(email="to")
             def test_email_alert(x):
                 raise RuntimeError("Fail please")
@@ -2005,10 +2057,10 @@ def test_outlet_datasets(self, create_task_instance):
         assert session.query(DatasetDagRunQueue.target_dag_id).filter_by(
             dataset_id=event.dataset.id
         ).order_by(DatasetDagRunQueue.target_dag_id).all() == [
-            ("dataset_consumes_1",),
-            ("dataset_consumes_1_and_2",),
-            ("dataset_consumes_1_never_scheduled",),
-        ]
+                   ("dataset_consumes_1",),
+                   ("dataset_consumes_1_and_2",),
+                   ("dataset_consumes_1_never_scheduled",),
+               ]
 
         # check that one event record created for dataset1 and this TI
         assert session.query(DatasetModel.uri).join(DatasetEvent.dataset).filter(
@@ -2113,7 +2165,6 @@ def test_changing_of_dataset_when_ddrq_is_already_populated(self, dag_maker, ses
         from airflow import Dataset
 
         with dag_maker(schedule=None, serialized=True) as dag1:
-
             @task(outlets=Dataset("test/1"))
             def test_task1():
                 print(1)
@@ -2124,7 +2175,6 @@ def test_task1():
         test_task1 = dag1.get_task("test_task1")
 
         with dag_maker(dag_id="testdag", schedule=[Dataset("test/1")], serialized=True):
-
             @task
             def test_task2():
                 print(1)
@@ -2135,7 +2185,6 @@ def test_task2():
         ti.run()
         # Change the dataset.
         with dag_maker(dag_id="testdag", schedule=[Dataset("test2/1")], serialized=True):
-
             @task
             def test_task2():
                 print(1)
@@ -3265,7 +3314,6 @@ def setup_class(self):
     def test_not_recorded_if_leaf(self, dag_maker, xcom_value):
         """Return value should not be recorded if there are no downstreams."""
         with dag_maker(dag_id="test_not_recorded_for_unused") as dag:
-
             @dag.task()
             def push_something():
                 return xcom_value
@@ -3281,7 +3329,6 @@ def push_something():
     def test_not_recorded_if_not_used(self, dag_maker, xcom_value):
         """Return value should not be recorded if no downstreams are mapped."""
         with dag_maker(dag_id="test_not_recorded_for_unused") as dag:
-
             @dag.task()
             def push_something():
                 return xcom_value
@@ -3302,7 +3349,6 @@ def completely_different():
     def test_not_recorded_if_irrelevant(self, dag_maker, xcom_1, xcom_4):
         """Return value should only be recorded if a mapped downstream uses the it."""
         with dag_maker(dag_id="test_not_recorded_for_unused") as dag:
-
             @dag.task()
             def push_1():
                 return xcom_1
@@ -3355,7 +3401,6 @@ def tg(arg):
     def test_expand_error_if_unmappable_type(self, dag_maker, return_value, exception_type, error_message):
         """If an unmappable return value is used for expand(), fail the task that pushed the XCom."""
         with dag_maker(dag_id="test_expand_error_if_unmappable_type") as dag:
-
             @dag.task()
             def push_something():
                 return return_value
@@ -3390,7 +3435,6 @@ def test_expand_kwargs_error_if_unmappable_type(
     ):
         """If an unmappable return value is used for expand_kwargs(), fail the task that pushed the XCom."""
         with dag_maker(dag_id="test_expand_kwargs_error_if_unmappable_type") as dag:
-
             @dag.task()
             def push():
                 return return_value
@@ -3421,7 +3465,6 @@ def test_task_group_expand_error_if_unmappable_type(
     ):
         """If an unmappable return value is used , fail the task that pushed the XCom."""
         with dag_maker(dag_id="test_task_group_expand_error_if_unmappable_type") as dag:
-
             @dag.task()
             def push():
                 return return_value
@@ -3456,7 +3499,6 @@ def test_task_group_expand_kwargs_error_if_unmappable_type(
     ):
         """If an unmappable return value is used, fail the task that pushed the XCom."""
         with dag_maker(dag_id="test_task_group_expand_kwargs_error_if_unmappable_type") as dag:
-
             @dag.task()
             def push():
                 return return_value
@@ -3568,7 +3610,6 @@ def pull(arg1, arg2):
     def test_error_if_upstream_does_not_push(self, dag_maker):
         """Fail the upstream task if it fails to push the XCom used for task mapping."""
         with dag_maker(dag_id="test_not_recorded_for_unused") as dag:
-
             @dag.task(do_xcom_push=False)
             def push_something():
                 return [1, 2]
@@ -3591,7 +3632,6 @@ def pull_something(value):
     def test_error_if_unmappable_length(self, dag_maker):
         """If an unmappable return value is used to map, fail the task that pushed the XCom."""
         with dag_maker(dag_id="test_not_recorded_for_unused") as dag:
-
             @dag.task()
             def push_something():
                 return [1, 2]
@@ -3620,7 +3660,6 @@ def pull_something(value):
     def test_written_task_map(self, dag_maker, xcom_value, expected_length, expected_keys):
         """Return value should be recorded in TaskMap if it's used by a downstream to map."""
         with dag_maker(dag_id="test_written_task_map") as dag:
-
             @dag.task()
             def push_something():
                 return xcom_value
@@ -3656,7 +3695,6 @@ def test_map_literal(self, literal, expected_outputs, dag_maker, session):
         outputs = []
 
         with dag_maker(dag_id="literal", session=session) as dag:
-
             @dag.task
             def show(value):
                 outputs.append(value)
@@ -3689,7 +3727,6 @@ def test_map_xcom(self, upstream_return, expected_outputs, dag_maker, session):
         outputs = []
 
         with dag_maker(dag_id="xcom", session=session) as dag:
-
             @dag.task
             def emit():
                 return upstream_return
@@ -3760,7 +3797,6 @@ def test_map_product_same(self, dag_maker, session):
         outputs = []
 
         with dag_maker(dag_id="product_same", session=session) as dag:
-
             @dag.task
             def emit_numbers():
                 return [1, 2]
@@ -3793,7 +3829,6 @@ def test_map_literal_cross_product(self, dag_maker, session):
         outputs = []
 
         with dag_maker(dag_id="product_same_types", session=session) as dag:
-
             @dag.task
             def show(a, b):
                 outputs.append((a, b))
@@ -4028,7 +4063,6 @@ class SimpleBashOperator(BashOperator):
         template_fields = ()
 
     with dag_maker(dag_id="product_same_types", session=session) as dag:
-
         @dag.task
         def get_extra_env():
             return [{"foo": "bar"}, {"foo": "biz"}]
@@ -4048,7 +4082,6 @@ def test_mapped_task_does_not_error_in_mini_scheduler_if_upstreams_are_not_done(
     not marked as `upstream_failed'
     """
     with dag_maker() as dag:
-
         @dag.task
         def second_task():
             return [0, 1, 2]
@@ -4090,7 +4123,6 @@ def test_empty_operator_is_not_considered_in_mini_scheduler(dag_maker, caplog, s
     submit them directly to worker.
     """
     with dag_maker() as dag:
-
         @dag.task
         def first_task():
             print(2)
@@ -4130,7 +4162,6 @@ def second_task():
 def test_mapped_task_expands_in_mini_scheduler_if_upstreams_are_done(dag_maker, caplog, session):
     """Test that mini scheduler expands mapped task"""
     with dag_maker() as dag:
-
         @dag.task
         def second_task():
             return [0, 1, 2]
@@ -4170,7 +4201,6 @@ def last_task():
 
 def test_mini_scheduler_not_skip_mapped_downstream_until_all_upstreams_finish(dag_maker, session):
     with dag_maker(session=session):
-
         @task
         def generate() -> list[list[int]]:
             return []

From a9c4573943708e213fc0b71e718d3de9ef0a3cfa Mon Sep 17 00:00:00 2001
From: "maciej.dembiczak" <maciej.dembiczak@openx.com>
Date: Tue, 30 May 2023 15:33:27 +0200
Subject: [PATCH 3/3] fix tests and code style

---
 tests/models/test_taskinstance.py | 115 +++++++++++++++++-------------
 1 file changed, 67 insertions(+), 48 deletions(-)

diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
index 9f62fa9833ca3..3ae696a27c142 100644
--- a/tests/models/test_taskinstance.py
+++ b/tests/models/test_taskinstance.py
@@ -82,7 +82,7 @@
 from airflow.utils.db import merge_conn
 from airflow.utils.module_loading import qualname
 from airflow.utils.session import create_session, provide_session
-from airflow.utils.state import State, TaskInstanceState, DagRunState
+from airflow.utils.state import DagRunState, State, TaskInstanceState
 from airflow.utils.task_group import TaskGroup
 from airflow.utils.types import DagRunType
 from airflow.utils.xcom import XCOM_RETURN_KEY
@@ -1773,55 +1773,51 @@ def test_get_num_running_task_instances(self, create_task_instance):
     def test_get_num_running_task_instances_per_dagrun(self, create_task_instance, dag_maker):
         session = settings.Session()
 
-        @task
-        def dummy_mapped_task(x: int):
-            return x
-
-        with dag_maker(dag_id="test_dag") as dag:
-            dummy_mapped_task.expand(x=[1, 2])
-            dagrun1 = dag.create_dagrun(execution_date=timezone.utcnow(),
-                                       state=DagRunState.RUNNING,
-                                       run_id="run_id_1",
-                                       session=session)
-
-            ti1 = dagrun1.task_instances[0]
-            ti2 = dagrun1.task_instances[1]
-
-        with dag_maker(dag_id="test_dag") as dag:
-            dummy_mapped_task.expand(x=[3])
-            dagrun2 = dag.create_dagrun(execution_date=timezone.utcnow(),
-                                       state=DagRunState.RUNNING,
-                                       run_id="run_id_2",
-                                       session=session)
-
-            ti3 = dagrun2.task_instances[0]
+        with dag_maker(dag_id="test_dag"):
+            MockOperator.partial(task_id="task_1").expand_kwargs([{"a": 1, "b": 2}, {"a": 3, "b": 4}])
+            MockOperator.partial(task_id="task_2").expand_kwargs([{"a": 1, "b": 2}])
+            MockOperator.partial(task_id="task_3").expand_kwargs([{"a": 1, "b": 2}])
 
-        with dag_maker(dag_id="test_dag") as dag:
-            dummy_mapped_task.expand(x=[4])
-            dagrun3 = dag.create_dagrun(execution_date=timezone.utcnow(),
-                                       state=DagRunState.RUNNING,
-                                       run_id="run_id_3",
-                                       session=session)
-
-            ti4 = dagrun3.task_instances[0]
-
-        assert ti1 in session
-        assert ti2 in session
-        assert ti3 in session
-        assert ti4 in session
+        dr1 = dag_maker.create_dagrun(
+            execution_date=timezone.utcnow(), state=DagRunState.RUNNING, run_id="run_id_1", session=session
+        )
+        tis1 = {(ti.task_id, ti.map_index): ti for ti in dr1.task_instances}
+        print(f"tis1: {tis1}")
 
-        ti1.state = State.RUNNING
-        ti2.state = State.QUEUED
-        ti3.state = State.RUNNING
-        ti4.state = State.RUNNING
+        dr2 = dag_maker.create_dagrun(
+            execution_date=timezone.utcnow(), state=DagRunState.RUNNING, run_id="run_id_2", session=session
+        )
+        tis2 = {(ti.task_id, ti.map_index): ti for ti in dr2.task_instances}
+
+        assert tis1[("task_1", 0)] in session
+        assert tis1[("task_1", 1)] in session
+        assert tis1[("task_2", 0)] in session
+        assert tis1[("task_3", 0)] in session
+        assert tis2[("task_1", 0)] in session
+        assert tis2[("task_1", 1)] in session
+        assert tis2[("task_2", 0)] in session
+        assert tis2[("task_3", 0)] in session
+
+        tis1[("task_1", 0)].state = State.RUNNING
+        tis1[("task_1", 1)].state = State.QUEUED
+        tis1[("task_2", 0)].state = State.RUNNING
+        tis1[("task_3", 0)].state = State.RUNNING
+        tis2[("task_1", 0)].state = State.RUNNING
+        tis2[("task_1", 1)].state = State.QUEUED
+        tis2[("task_2", 0)].state = State.RUNNING
+        tis2[("task_3", 0)].state = State.RUNNING
 
         session.commit()
 
-        assert 1 == ti1.get_num_running_task_instances(session=session, same_dagrun=True)
-        assert 1 == ti2.get_num_running_task_instances(session=session, same_dagrun=True)
-        assert 3 == ti3.get_num_running_task_instances(session=session)
-        assert 1 == ti4.get_num_running_task_instances(session=session, same_dagrun=True)
+        assert 1 == tis1[("task_1", 0)].get_num_running_task_instances(session=session, same_dagrun=True)
+        assert 1 == tis1[("task_1", 1)].get_num_running_task_instances(session=session, same_dagrun=True)
+        assert 2 == tis1[("task_2", 0)].get_num_running_task_instances(session=session)
+        assert 1 == tis1[("task_3", 0)].get_num_running_task_instances(session=session, same_dagrun=True)
 
+        assert 1 == tis2[("task_1", 0)].get_num_running_task_instances(session=session, same_dagrun=True)
+        assert 1 == tis2[("task_1", 1)].get_num_running_task_instances(session=session, same_dagrun=True)
+        assert 2 == tis2[("task_2", 0)].get_num_running_task_instances(session=session)
+        assert 1 == tis2[("task_3", 0)].get_num_running_task_instances(session=session, same_dagrun=True)
 
     def test_log_url(self, create_task_instance):
         ti = create_task_instance(dag_id="dag", task_id="op", execution_date=timezone.datetime(2018, 1, 1))
@@ -1956,6 +1952,7 @@ def test_email_alert_with_filenotfound_config(self, mock_send_email, dag_maker):
     @patch("airflow.models.taskinstance.send_email")
     def test_failure_mapped_taskflow(self, mock_send_email, dag_maker, session, task_id):
         with dag_maker(session=session) as dag:
+
             @dag.task(email="to")
             def test_email_alert(x):
                 raise RuntimeError("Fail please")
@@ -2057,10 +2054,10 @@ def test_outlet_datasets(self, create_task_instance):
         assert session.query(DatasetDagRunQueue.target_dag_id).filter_by(
             dataset_id=event.dataset.id
         ).order_by(DatasetDagRunQueue.target_dag_id).all() == [
-                   ("dataset_consumes_1",),
-                   ("dataset_consumes_1_and_2",),
-                   ("dataset_consumes_1_never_scheduled",),
-               ]
+            ("dataset_consumes_1",),
+            ("dataset_consumes_1_and_2",),
+            ("dataset_consumes_1_never_scheduled",),
+        ]
 
         # check that one event record created for dataset1 and this TI
         assert session.query(DatasetModel.uri).join(DatasetEvent.dataset).filter(
@@ -2165,6 +2162,7 @@ def test_changing_of_dataset_when_ddrq_is_already_populated(self, dag_maker, ses
         from airflow import Dataset
 
         with dag_maker(schedule=None, serialized=True) as dag1:
+
             @task(outlets=Dataset("test/1"))
             def test_task1():
                 print(1)
@@ -2175,6 +2173,7 @@ def test_task1():
         test_task1 = dag1.get_task("test_task1")
 
         with dag_maker(dag_id="testdag", schedule=[Dataset("test/1")], serialized=True):
+
             @task
             def test_task2():
                 print(1)
@@ -2185,6 +2184,7 @@ def test_task2():
         ti.run()
         # Change the dataset.
         with dag_maker(dag_id="testdag", schedule=[Dataset("test2/1")], serialized=True):
+
             @task
             def test_task2():
                 print(1)
@@ -3314,6 +3314,7 @@ def setup_class(self):
     def test_not_recorded_if_leaf(self, dag_maker, xcom_value):
         """Return value should not be recorded if there are no downstreams."""
         with dag_maker(dag_id="test_not_recorded_for_unused") as dag:
+
             @dag.task()
             def push_something():
                 return xcom_value
@@ -3329,6 +3330,7 @@ def push_something():
     def test_not_recorded_if_not_used(self, dag_maker, xcom_value):
         """Return value should not be recorded if no downstreams are mapped."""
         with dag_maker(dag_id="test_not_recorded_for_unused") as dag:
+
             @dag.task()
             def push_something():
                 return xcom_value
@@ -3349,6 +3351,7 @@ def completely_different():
     def test_not_recorded_if_irrelevant(self, dag_maker, xcom_1, xcom_4):
         """Return value should only be recorded if a mapped downstream uses the it."""
         with dag_maker(dag_id="test_not_recorded_for_unused") as dag:
+
             @dag.task()
             def push_1():
                 return xcom_1
@@ -3401,6 +3404,7 @@ def tg(arg):
     def test_expand_error_if_unmappable_type(self, dag_maker, return_value, exception_type, error_message):
         """If an unmappable return value is used for expand(), fail the task that pushed the XCom."""
         with dag_maker(dag_id="test_expand_error_if_unmappable_type") as dag:
+
             @dag.task()
             def push_something():
                 return return_value
@@ -3435,6 +3439,7 @@ def test_expand_kwargs_error_if_unmappable_type(
     ):
         """If an unmappable return value is used for expand_kwargs(), fail the task that pushed the XCom."""
         with dag_maker(dag_id="test_expand_kwargs_error_if_unmappable_type") as dag:
+
             @dag.task()
             def push():
                 return return_value
@@ -3465,6 +3470,7 @@ def test_task_group_expand_error_if_unmappable_type(
     ):
         """If an unmappable return value is used , fail the task that pushed the XCom."""
         with dag_maker(dag_id="test_task_group_expand_error_if_unmappable_type") as dag:
+
             @dag.task()
             def push():
                 return return_value
@@ -3499,6 +3505,7 @@ def test_task_group_expand_kwargs_error_if_unmappable_type(
     ):
         """If an unmappable return value is used, fail the task that pushed the XCom."""
         with dag_maker(dag_id="test_task_group_expand_kwargs_error_if_unmappable_type") as dag:
+
             @dag.task()
             def push():
                 return return_value
@@ -3610,6 +3617,7 @@ def pull(arg1, arg2):
     def test_error_if_upstream_does_not_push(self, dag_maker):
         """Fail the upstream task if it fails to push the XCom used for task mapping."""
         with dag_maker(dag_id="test_not_recorded_for_unused") as dag:
+
             @dag.task(do_xcom_push=False)
             def push_something():
                 return [1, 2]
@@ -3632,6 +3640,7 @@ def pull_something(value):
     def test_error_if_unmappable_length(self, dag_maker):
         """If an unmappable return value is used to map, fail the task that pushed the XCom."""
         with dag_maker(dag_id="test_not_recorded_for_unused") as dag:
+
             @dag.task()
             def push_something():
                 return [1, 2]
@@ -3660,6 +3669,7 @@ def pull_something(value):
     def test_written_task_map(self, dag_maker, xcom_value, expected_length, expected_keys):
         """Return value should be recorded in TaskMap if it's used by a downstream to map."""
         with dag_maker(dag_id="test_written_task_map") as dag:
+
             @dag.task()
             def push_something():
                 return xcom_value
@@ -3695,6 +3705,7 @@ def test_map_literal(self, literal, expected_outputs, dag_maker, session):
         outputs = []
 
         with dag_maker(dag_id="literal", session=session) as dag:
+
             @dag.task
             def show(value):
                 outputs.append(value)
@@ -3727,6 +3738,7 @@ def test_map_xcom(self, upstream_return, expected_outputs, dag_maker, session):
         outputs = []
 
         with dag_maker(dag_id="xcom", session=session) as dag:
+
             @dag.task
             def emit():
                 return upstream_return
@@ -3797,6 +3809,7 @@ def test_map_product_same(self, dag_maker, session):
         outputs = []
 
         with dag_maker(dag_id="product_same", session=session) as dag:
+
             @dag.task
             def emit_numbers():
                 return [1, 2]
@@ -3829,6 +3842,7 @@ def test_map_literal_cross_product(self, dag_maker, session):
         outputs = []
 
         with dag_maker(dag_id="product_same_types", session=session) as dag:
+
             @dag.task
             def show(a, b):
                 outputs.append((a, b))
@@ -4063,6 +4077,7 @@ class SimpleBashOperator(BashOperator):
         template_fields = ()
 
     with dag_maker(dag_id="product_same_types", session=session) as dag:
+
         @dag.task
         def get_extra_env():
             return [{"foo": "bar"}, {"foo": "biz"}]
@@ -4082,6 +4097,7 @@ def test_mapped_task_does_not_error_in_mini_scheduler_if_upstreams_are_not_done(
     not marked as `upstream_failed'
     """
     with dag_maker() as dag:
+
         @dag.task
         def second_task():
             return [0, 1, 2]
@@ -4123,6 +4139,7 @@ def test_empty_operator_is_not_considered_in_mini_scheduler(dag_maker, caplog, s
     submit them directly to worker.
     """
     with dag_maker() as dag:
+
         @dag.task
         def first_task():
             print(2)
@@ -4162,6 +4179,7 @@ def second_task():
 def test_mapped_task_expands_in_mini_scheduler_if_upstreams_are_done(dag_maker, caplog, session):
     """Test that mini scheduler expands mapped task"""
     with dag_maker() as dag:
+
         @dag.task
         def second_task():
             return [0, 1, 2]
@@ -4201,6 +4219,7 @@ def last_task():
 
 def test_mini_scheduler_not_skip_mapped_downstream_until_all_upstreams_finish(dag_maker, session):
     with dag_maker(session=session):
+
         @task
         def generate() -> list[list[int]]:
             return []
