Use airflow.datasets.Dataset in examples and tests

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]

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
 

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)

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"

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

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:
 

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