Use airflow.datasets.Dataset in examples and tests (apache#34605)

asapekia · Sep 25, 2023 · 5ac1e7c · 5ac1e7c
1 parent d0ae60f
commit 5ac1e7c
Show file tree

Hide file tree

Showing 7 changed files with 13 additions and 8 deletions.
diff --git a/airflow/example_dags/example_datasets.py b/airflow/example_dags/example_datasets.py
@@ -40,7 +40,8 @@
 
 import pendulum
 
-from airflow import DAG, Dataset
+from airflow.datasets import Dataset
+from airflow.models.dag import DAG
 from airflow.operators.bash import BashOperator
 
 # [START dataset_def]

diff --git a/tests/dags/test_datasets.py b/tests/dags/test_datasets.py
@@ -19,8 +19,9 @@
 
 from datetime import datetime
 
-from airflow import DAG, Dataset
+from airflow.datasets import Dataset
 from airflow.exceptions import AirflowFailException, AirflowSkipException
+from airflow.models.dag import DAG
 from airflow.operators.bash import BashOperator
 from airflow.operators.python import PythonOperator
 

diff --git a/tests/dags/test_only_empty_tasks.py b/tests/dags/test_only_empty_tasks.py
@@ -20,7 +20,8 @@
 from datetime import datetime
 from typing import Sequence
 
-from airflow import DAG, Dataset
+from airflow.datasets import Dataset
+from airflow.models.dag import DAG
 from airflow.operators.empty import EmptyOperator
 
 DEFAULT_DATE = datetime(2016, 1, 1)

diff --git a/tests/decorators/test_python.py b/tests/decorators/test_python.py
@@ -893,7 +893,7 @@ def other(x):
 
 
 def test_task_decorator_dataset(dag_maker, session):
-    from airflow import Dataset
+    from airflow.datasets import Dataset
 
     result = None
     uri = "s3://test"

diff --git a/tests/models/test_serialized_dag.py b/tests/models/test_serialized_dag.py
@@ -23,8 +23,10 @@
 import pendulum
 import pytest
 
-from airflow import DAG, Dataset, example_dags as example_dags_module
-from airflow.models import DagBag
+import airflow.example_dags as example_dags_module
+from airflow.datasets import Dataset
+from airflow.models.dag import DAG
+from airflow.models.dagbag import DagBag
 from airflow.models.dagcode import DagCode
 from airflow.models.serialized_dag import SerializedDagModel as SDM
 from airflow.operators.bash import BashOperator

diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
@@ -2196,7 +2196,7 @@ def test_changing_of_dataset_when_ddrq_is_already_populated(self, dag_maker, ses
         Test that when a task that produces dataset has ran, that changing the consumer
         dag dataset will not cause primary key blank-out
         """
-        from airflow import Dataset
+        from airflow.datasets import Dataset
 
         with dag_maker(schedule=None, serialized=True) as dag1:
 

diff --git a/tests/www/views/test_views_dataset.py b/tests/www/views/test_views_dataset.py
@@ -21,7 +21,7 @@
 import pytest
 from dateutil.tz import UTC
 
-from airflow import Dataset
+from airflow.datasets import Dataset
 from airflow.models.dataset import DatasetEvent, DatasetModel
 from airflow.operators.empty import EmptyOperator
 from tests.test_utils.asserts import assert_queries_count