[airflow] branch main updated: Use `airflow.datasets.Dataset` in examples and tests (#34605)

[hu...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

husseinawala pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new 5ac1e7c355 Use `airflow.datasets.Dataset` in examples and tests (#34605)
5ac1e7c355 is described below

commit 5ac1e7c355c83346167c6a485290a5a86fcdfb5b
Author: Andrey Anshin <An...@taragol.is>
AuthorDate: Mon Sep 25 22:07:58 2023 +0400

    Use `airflow.datasets.Dataset` in examples and tests (#34605)
---
 airflow/example_dags/example_datasets.py | 3 ++-
 tests/dags/test_datasets.py              | 3 ++-
 tests/dags/test_only_empty_tasks.py      | 3 ++-
 tests/decorators/test_python.py          | 2 +-
 tests/models/test_serialized_dag.py      | 6 ++++--
 tests/models/test_taskinstance.py        | 2 +-
 tests/www/views/test_views_dataset.py    | 2 +-
 7 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/airflow/example_dags/example_datasets.py b/airflow/example_dags/example_datasets.py
index c2db9b1968..ba3ed927f1 100644
--- a/airflow/example_dags/example_datasets.py
+++ b/airflow/example_dags/example_datasets.py
@@ -40,7 +40,8 @@ from __future__ import annotations
 
 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
index 70dd65d558..4bdef9f697 100644
--- a/tests/dags/test_datasets.py
+++ b/tests/dags/test_datasets.py
@@ -19,8 +19,9 @@ from __future__ import annotations
 
 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
index 66e31ce38e..68f1dc5e89 100644
--- a/tests/dags/test_only_empty_tasks.py
+++ b/tests/dags/test_only_empty_tasks.py
@@ -20,7 +20,8 @@ from __future__ import annotations
 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
index 1368ddb434..a4cd325a1f 100644
--- a/tests/decorators/test_python.py
+++ b/tests/decorators/test_python.py
@@ -893,7 +893,7 @@ def test_no_warnings(reset_logging_config, caplog):
 
 
 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
index 61be12676f..eec8ec7cbf 100644
--- a/tests/models/test_serialized_dag.py
+++ b/tests/models/test_serialized_dag.py
@@ -23,8 +23,10 @@ from unittest import mock
 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
index d43f7be645..db863d6dbd 100644
--- a/tests/models/test_taskinstance.py
+++ b/tests/models/test_taskinstance.py
@@ -2196,7 +2196,7 @@ class TestTaskInstance:
         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
index f2968faea1..ae08d94a3e 100644
--- a/tests/www/views/test_views_dataset.py
+++ b/tests/www/views/test_views_dataset.py
@@ -21,7 +21,7 @@ import pendulum
 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