Adjust test for Airflow-2.9

astronomer · May 6, 2024 · 75ef5f4 · 75ef5f4
1 parent 0960c84
commit 75ef5f4
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 8 deletions.
diff --git a/python-sdk/tests/airflow_tests/test_datasets.py b/python-sdk/tests/airflow_tests/test_datasets.py
@@ -104,7 +104,6 @@ def test_kwargs_with_temp_table():
 @pytest.mark.skipif(airflow.__version__ < "2.4.0", reason="Require Airflow version >= 2.4.0")
 def test_example_dataset_dag():
     from airflow.datasets import Dataset
-    from airflow.models.dataset import DatasetModel
 
     dir_path = os.path.dirname(os.path.realpath(__file__))
     db = DagBag(dir_path + "/../../example_dags/example_datasets.py")
@@ -115,9 +114,8 @@ def test_example_dataset_dag():
     outlets = producer_dag.tasks[-1].outlets
     assert isinstance(outlets[0], Dataset)
     # Test that dataset_triggers is only set if all the instances passed to the DAG object are Datasets
-    assert consumer_dag.dataset_triggers == outlets
+    assert consumer_dag.dataset_triggers.objects[0] == outlets[0]
     assert outlets[0].uri == "astro://postgres_conn@?table=imdb_movies"
-    assert DatasetModel.from_public(outlets[0]) == Dataset("astro://postgres_conn@?table=imdb_movies")
 
 
 def test_disable_auto_inlets_outlets():

diff --git a/python-sdk/tests/sql/operators/test_cleanup.py b/python-sdk/tests/sql/operators/test_cleanup.py
@@ -105,12 +105,12 @@ def test_error_raised_with_blocking_op_executors(
     reason="BackfillJobRunner and Job classes are only available in airflow >= 2.6",
 )
 @pytest.mark.parametrize(
-    "executor_in_job,executor_in_cfg,expected_val",
+    "executor_in_job, executor_in_cfg, expected_val",
     [
         (SequentialExecutor(), "LocalExecutor", True),
         (LocalExecutor(), "LocalExecutor", False),
-        (None, "LocalExecutor", False),
-        (None, "SequentialExecutor", True),
+        (LocalExecutor(), "LocalExecutor", False),
+        (SequentialExecutor(), "SequentialExecutor", True),
     ],
 )
 def test_single_worker_mode_backfill(executor_in_job, executor_in_cfg, expected_val):
@@ -177,8 +177,8 @@ def test_single_worker_mode_backfill_airflow_2_5(executor_in_job, executor_in_cf
     [
         (SequentialExecutor(), "LocalExecutor", True),
         (LocalExecutor(), "LocalExecutor", False),
-        (None, "LocalExecutor", False),
-        (None, "SequentialExecutor", True),
+        (LocalExecutor(), "LocalExecutor", False),
+        (SequentialExecutor(), "SequentialExecutor", True),
     ],
 )
 def test_single_worker_mode_scheduler_job(executor_in_job, executor_in_cfg, expected_val):