fix more tests

Author: dstandish

providers/standard/tests/unit/standard/operators/test_python.py

@@ -1915,7 +1915,7 @@ class TestShortCircuitWithTeardown:
     def test_short_circuit_with_teardowns(
         self, dag_maker, ignore_downstream_trigger_rules, should_skip, with_teardown, expected
     ):
-        with dag_maker() as dag:
+        with dag_maker(serialized=True):
             op1 = ShortCircuitOperator(
                 task_id="op1",
                 python_callable=lambda: not should_skip,
@@ -1928,21 +1928,20 @@ def test_short_circuit_with_teardowns(
                 op4.as_teardown()
             op1 >> op2 >> op3 >> op4
             op1.skip = MagicMock()
-            dagrun = dag_maker.create_dagrun()
-            tis = dagrun.get_task_instances()
-            ti: TaskInstance = next(x for x in tis if x.task_id == "op1")
-            ti._run_raw_task()
-            expected_tasks = {dag.task_dict[x] for x in expected}
+        dagrun = dag_maker.create_dagrun()
+        tis = dagrun.get_task_instances()
+        ti: TaskInstance = next(x for x in tis if x.task_id == "op1")
+        ti._run_raw_task()
         if should_skip:
             # we can't use assert_called_with because it's a set and therefore not ordered
-            actual_skipped = set(op1.skip.call_args.kwargs["tasks"])
-            assert actual_skipped == expected_tasks
+            actual_skipped = set(x.task_id for x in op1.skip.call_args.kwargs["tasks"])
+            assert actual_skipped == set(expected)
         else:
             op1.skip.assert_not_called()
 
     @pytest.mark.parametrize("config", ["sequence", "parallel"])
     def test_short_circuit_with_teardowns_complicated(self, dag_maker, config):
-        with dag_maker():
+        with dag_maker(serialized=True):
             s1 = PythonOperator(task_id="s1", python_callable=print).as_setup()
             s2 = PythonOperator(task_id="s2", python_callable=print).as_setup()
             op1 = ShortCircuitOperator(
@@ -1959,16 +1958,16 @@ def test_short_circuit_with_teardowns_complicated(self, dag_maker, config):
             else:
                 raise ValueError("unexpected")
             op1.skip = MagicMock()
-            dagrun = dag_maker.create_dagrun()
-            tis = dagrun.get_task_instances()
-            ti: TaskInstance = next(x for x in tis if x.task_id == "op1")
-            ti._run_raw_task()
-            # we can't use assert_called_with because it's a set and therefore not ordered
-            actual_skipped = set(op1.skip.call_args.kwargs["tasks"])
-            assert actual_skipped == {s2, op2}
+        dagrun = dag_maker.create_dagrun()
+        tis = dagrun.get_task_instances()
+        ti: TaskInstance = next(x for x in tis if x.task_id == "op1")
+        ti._run_raw_task()
+        # we can't use assert_called_with because it's a set and therefore not ordered
+        actual_skipped = set(op1.skip.call_args.kwargs["tasks"])
+        assert actual_skipped == {s2, op2}
 
     def test_short_circuit_with_teardowns_complicated_2(self, dag_maker):
-        with dag_maker():
+        with dag_maker(serialized=True):
             s1 = PythonOperator(task_id="s1", python_callable=print).as_setup()
             s2 = PythonOperator(task_id="s2", python_callable=print).as_setup()
             op1 = ShortCircuitOperator(
@@ -1986,22 +1985,22 @@ def test_short_circuit_with_teardowns_complicated_2(self, dag_maker):
             # in this case we don't want to skip t2 since it should run
             op1 >> t2
             op1.skip = MagicMock()
-            dagrun = dag_maker.create_dagrun()
-            tis = dagrun.get_task_instances()
-            ti: TaskInstance = next(x for x in tis if x.task_id == "op1")
-            ti._run_raw_task()
-            # we can't use assert_called_with because it's a set and therefore not ordered
-            actual_kwargs = op1.skip.call_args.kwargs
-            actual_skipped = set(actual_kwargs["tasks"])
-            assert actual_skipped == {op3}
+        dagrun = dag_maker.create_dagrun()
+        tis = dagrun.get_task_instances()
+        ti: TaskInstance = next(x for x in tis if x.task_id == "op1")
+        ti._run_raw_task()
+        # we can't use assert_called_with because it's a set and therefore not ordered
+        actual_kwargs = op1.skip.call_args.kwargs
+        actual_skipped = set(actual_kwargs["tasks"])
+        assert actual_skipped == {op3}
 
     @pytest.mark.parametrize("level", [logging.DEBUG, logging.INFO])
     def test_short_circuit_with_teardowns_debug_level(self, dag_maker, level, clear_db):
         """
         When logging is debug we convert to a list to log the tasks skipped
         before passing them to the skip method.
         """
-        with dag_maker():
+        with dag_maker(serialized=True):
             s1 = PythonOperator(task_id="s1", python_callable=print).as_setup()
             s2 = PythonOperator(task_id="s2", python_callable=print).as_setup()
             op1 = ShortCircuitOperator(
@@ -2020,18 +2019,18 @@ def test_short_circuit_with_teardowns_debug_level(self, dag_maker, level, clear_
             # in this case we don't want to skip t2 since it should run
             op1 >> t2
             op1.skip = MagicMock()
-            dagrun = dag_maker.create_dagrun()
-            tis = dagrun.get_task_instances()
-            ti: TaskInstance = next(x for x in tis if x.task_id == "op1")
-            ti._run_raw_task()
-            # we can't use assert_called_with because it's a set and therefore not ordered
-            actual_kwargs = op1.skip.call_args.kwargs
-            actual_skipped = actual_kwargs["tasks"]
-            if level <= logging.DEBUG:
-                assert isinstance(actual_skipped, list)
-            else:
-                assert isinstance(actual_skipped, Generator)
-            assert set(actual_skipped) == {op3}
+        dagrun = dag_maker.create_dagrun()
+        tis = dagrun.get_task_instances()
+        ti: TaskInstance = next(x for x in tis if x.task_id == "op1")
+        ti._run_raw_task()
+        # we can't use assert_called_with because it's a set and therefore not ordered
+        actual_kwargs = op1.skip.call_args.kwargs
+        actual_skipped = actual_kwargs["tasks"]
+        if level <= logging.DEBUG:
+            assert isinstance(actual_skipped, list)
+        else:
+            assert isinstance(actual_skipped, Generator)
+        assert set(actual_skipped) == {op3}
 
 
 @pytest.mark.parametrize(

providers/standard/tests/unit/standard/sensors/test_external_task_sensor.py

@@ -30,6 +30,7 @@
 from airflow.models import DagBag, DagRun, TaskInstance
 from airflow.models.baseoperator import BaseOperator
 from airflow.models.dag import DAG
+from airflow.models.dagbundle import DagBundleModel
 from airflow.models.serialized_dag import SerializedDagModel
 from airflow.models.xcom_arg import XComArg
 from airflow.providers.standard.operators.bash import BashOperator
@@ -1771,14 +1772,13 @@ def test_external_task_marker_cyclic_shallow(dag_bag_cyclic):
 
 
 @pytest.fixture
-def dag_bag_multiple():
+def dag_bag_multiple(session):
     """
     Create a DagBag containing two DAGs, linked by multiple ExternalTaskMarker.
     """
     dag_bag = DagBag(dag_folder=DEV_NULL, include_examples=False)
     daily_dag = DAG("daily_dag", start_date=DEFAULT_DATE, schedule="@daily")
     agg_dag = DAG("agg_dag", start_date=DEFAULT_DATE, schedule="@daily")
-
     if AIRFLOW_V_3_0_PLUS:
         dag_bag.bag_dag(dag=daily_dag)
         dag_bag.bag_dag(dag=agg_dag)
@@ -1798,6 +1798,12 @@ def dag_bag_multiple():
             dag=agg_dag,
         )
         begin >> task
+    bundle_name = "abcbunhdlerch3rc"
+    session.merge(DagBundleModel(name=bundle_name))
+    session.commit()
+    DAG.bulk_write_to_db(bundle_name=bundle_name, dags=[daily_dag, agg_dag], bundle_version=None)
+    SerializedDagModel.write_dag(dag=daily_dag, bundle_name=bundle_name)
+    SerializedDagModel.write_dag(dag=agg_dag, bundle_name=bundle_name)
 
     return dag_bag
 
@@ -1819,7 +1825,7 @@ def test_clear_multiple_external_task_marker(dag_bag_multiple):
 
 
 @pytest.fixture
-def dag_bag_head_tail():
+def dag_bag_head_tail(session):
     """
     Create a DagBag containing one DAG, with task "head" depending on task "tail" of the
     previous logical_date.
@@ -1858,6 +1864,11 @@ def dag_bag_head_tail():
         dag_bag.bag_dag(dag=dag)
     else:
         dag_bag.bag_dag(dag=dag, root_dag=dag)
+    bundle_name = "9e8uh9odhu9c"
+    session.merge(DagBundleModel(name=bundle_name))
+    session.commit()
+    DAG.bulk_write_to_db(bundle_name=bundle_name, dags=[dag], bundle_version=None)
+    SerializedDagModel.write_dag(dag=dag, bundle_name=bundle_name)
 
     return dag_bag