diff --git a/airflow-core/tests/unit/cli/commands/test_dag_command.py b/airflow-core/tests/unit/cli/commands/test_dag_command.py index 2c16be86af281..e13cc3a410e73 100644 --- a/airflow-core/tests/unit/cli/commands/test_dag_command.py +++ b/airflow-core/tests/unit/cli/commands/test_dag_command.py @@ -38,7 +38,6 @@ from airflow.cli.commands import dag_command from airflow.exceptions import AirflowException from airflow.models import DagBag, DagModel, DagRun -from airflow.models.baseoperator import BaseOperator from airflow.models.serialized_dag import SerializedDagModel from airflow.providers.standard.triggers.temporal import DateTimeTrigger, TimeDeltaTrigger from airflow.sdk import task @@ -57,6 +56,11 @@ ) from unit.models import TEST_DAGS_FOLDER +try: + from airflow.sdk import BaseOperator +except ImportError: + from airflow.models.baseoperator import BaseOperator # type: ignore[no-redef] + DEFAULT_DATE = timezone.make_aware(datetime(2015, 1, 1), timezone=timezone.utc) if pendulum.__version__.startswith("3"): DEFAULT_DATE_REPR = DEFAULT_DATE.isoformat(sep=" ") diff --git a/airflow-core/tests/unit/cluster_policies/__init__.py b/airflow-core/tests/unit/cluster_policies/__init__.py index 88b1d329b70ed..e9c380838e90d 100644 --- a/airflow-core/tests/unit/cluster_policies/__init__.py +++ b/airflow-core/tests/unit/cluster_policies/__init__.py @@ -24,7 +24,7 @@ from airflow.configuration import conf from airflow.exceptions import AirflowClusterPolicySkipDag, AirflowClusterPolicyViolation -from airflow.sdk.bases.operator import BaseOperator +from airflow.sdk import BaseOperator if TYPE_CHECKING: from airflow.models.dag import DAG diff --git a/devel-common/src/tests_common/test_utils/mock_operators.py b/devel-common/src/tests_common/test_utils/mock_operators.py index d4fd43421f6d3..6082521b19066 100644 --- a/devel-common/src/tests_common/test_utils/mock_operators.py +++ b/devel-common/src/tests_common/test_utils/mock_operators.py @@ -21,17 +21,17 @@ import attr -from airflow.models.baseoperator import BaseOperator - from tests_common.test_utils.compat import BaseOperatorLink from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS if TYPE_CHECKING: from airflow.sdk.definitions.context import Context -if AIRFLOW_V_3_0_PLUS: +try: from airflow.models.xcom import XComModel as XCom -else: + from airflow.sdk import BaseOperator +except ImportError: + from airflow.models.baseoperator import BaseOperator # type: ignore[no-redef] from airflow.models.xcom import XCom # type: ignore[no-redef] diff --git a/devel-common/src/tests_common/test_utils/version_compat.py b/devel-common/src/tests_common/test_utils/version_compat.py index 2e4214f280a89..50d50d2a1cc7e 100644 --- a/devel-common/src/tests_common/test_utils/version_compat.py +++ b/devel-common/src/tests_common/test_utils/version_compat.py @@ -37,12 +37,16 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: AIRFLOW_V_3_0_3_PLUS = get_base_airflow_version_tuple() >= (3, 0, 3) AIRFLOW_V_3_1_PLUS = get_base_airflow_version_tuple() >= (3, 1, 0) + if AIRFLOW_V_3_1_PLUS: + from airflow.sdk import PokeReturnValue, timezone from airflow.sdk.bases.xcom import BaseXCom from airflow.sdk.definitions._internal.decorators import remove_task_decorator XCOM_RETURN_KEY = BaseXCom.XCOM_RETURN_KEY else: + from airflow.sensors.base import PokeReturnValue # type: ignore[no-redef] + from airflow.utils import timezone # type: ignore[attr-defined,no-redef] from airflow.utils.decorators import remove_task_decorator # type: ignore[no-redef] from airflow.utils.xcom import XCOM_RETURN_KEY # type: ignore[no-redef] @@ -66,5 +70,7 @@ def get_sqlalchemy_version_tuple() -> tuple[int, int, int]: "SQLALCHEMY_V_1_4", "SQLALCHEMY_V_2_0", "XCOM_RETURN_KEY", + "PokeReturnValue", "remove_task_decorator", + "timezone", ] diff --git a/performance/tests/test_performance_dag.py b/performance/tests/test_performance_dag.py index 387a9faf8c481..6cc72215d9750 100644 --- a/performance/tests/test_performance_dag.py +++ b/performance/tests/test_performance_dag.py @@ -22,9 +22,10 @@ import json import os +import re import pytest -import re + from airflow.configuration import conf from airflow.models import DagBag from airflow.utils.trigger_rule import TriggerRule diff --git a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py index 2316cd3c13dbc..a3dd26d686bca 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py +++ b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py @@ -70,8 +70,7 @@ from airflow.models import TaskInstance from airflow.providers.common.compat.assets import Asset - from airflow.sdk import DAG - from airflow.sdk.bases.operator import BaseOperator + from airflow.sdk import DAG, BaseOperator from airflow.sdk.definitions.mappedoperator import MappedOperator from airflow.sdk.execution_time.secrets_masker import ( Redactable, @@ -83,8 +82,7 @@ from airflow.utils.state import DagRunState, TaskInstanceState else: try: - from airflow.sdk import DAG - from airflow.sdk.bases.operator import BaseOperator + from airflow.sdk import DAG, BaseOperator from airflow.sdk.definitions.mappedoperator import MappedOperator except ImportError: from airflow.models import DAG, BaseOperator, MappedOperator diff --git a/providers/openlineage/tests/unit/openlineage/extractors/test_manager.py b/providers/openlineage/tests/unit/openlineage/extractors/test_manager.py index 100ac520d3d84..f18f7e22a9ef6 100644 --- a/providers/openlineage/tests/unit/openlineage/extractors/test_manager.py +++ b/providers/openlineage/tests/unit/openlineage/extractors/test_manager.py @@ -69,9 +69,8 @@ def hook_lineage_collector(): if AIRFLOW_V_3_0_PLUS: - from airflow.sdk import ObjectStoragePath + from airflow.sdk import BaseOperator, ObjectStoragePath from airflow.sdk.api.datamodels._generated import TaskInstance as SDKTaskInstance - from airflow.sdk.bases.operator import BaseOperator from airflow.sdk.execution_time import task_runner from airflow.sdk.execution_time.comms import StartupDetails from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance, parse diff --git a/providers/openlineage/tests/unit/openlineage/plugins/test_macros.py b/providers/openlineage/tests/unit/openlineage/plugins/test_macros.py index 2bfa74eb27fcd..a0b090ec53813 100644 --- a/providers/openlineage/tests/unit/openlineage/plugins/test_macros.py +++ b/providers/openlineage/tests/unit/openlineage/plugins/test_macros.py @@ -114,7 +114,7 @@ def test_lineage_parent_id(mock_run_id): @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Test only for Airflow 3.0+") def test_lineage_root_run_id_with_runtime_task_instance(create_runtime_ti): """Test lineage_root_run_id with real RuntimeTaskInstance object doesn't throw AttributeError.""" - from airflow.sdk.bases.operator import BaseOperator + from airflow.sdk import BaseOperator task = BaseOperator(task_id="test_task") diff --git a/providers/standard/src/airflow/providers/standard/example_dags/example_external_task_parent_deferrable.py b/providers/standard/src/airflow/providers/standard/example_dags/example_external_task_parent_deferrable.py index ffc04fabdf5ac..5b55fe566d722 100644 --- a/providers/standard/src/airflow/providers/standard/example_dags/example_external_task_parent_deferrable.py +++ b/providers/standard/src/airflow/providers/standard/example_dags/example_external_task_parent_deferrable.py @@ -20,7 +20,11 @@ from airflow.providers.standard.operators.empty import EmptyOperator from airflow.providers.standard.operators.trigger_dagrun import TriggerDagRunOperator from airflow.providers.standard.sensors.external_task import ExternalTaskSensor -from airflow.utils.timezone import datetime + +try: + from airflow.sdk.timezone import datetime +except ImportError: + from airflow.utils.timezone import datetime # type: ignore[no-redef] with DAG( dag_id="example_external_task", diff --git a/providers/standard/src/airflow/providers/standard/sensors/time.py b/providers/standard/src/airflow/providers/standard/sensors/time.py index 6d72c29fb4a81..0af548733da02 100644 --- a/providers/standard/src/airflow/providers/standard/sensors/time.py +++ b/providers/standard/src/airflow/providers/standard/sensors/time.py @@ -42,7 +42,10 @@ class StartTriggerArgs: # type: ignore[no-redef] timeout: datetime.timedelta | None = None -from airflow.utils import timezone +try: + from airflow.sdk import timezone +except ImportError: + from airflow.utils import timezone # type: ignore[attr-defined,no-redef] if TYPE_CHECKING: try: diff --git a/providers/standard/tests/unit/standard/sensors/test_date_time.py b/providers/standard/tests/unit/standard/sensors/test_date_time.py index 188237a957bd4..48be29f8ef891 100644 --- a/providers/standard/tests/unit/standard/sensors/test_date_time.py +++ b/providers/standard/tests/unit/standard/sensors/test_date_time.py @@ -25,7 +25,8 @@ from airflow import macros from airflow.models.dag import DAG from airflow.providers.standard.sensors.date_time import DateTimeSensor -from airflow.utils import timezone + +from tests_common.test_utils.version_compat import timezone DEFAULT_DATE = timezone.datetime(2015, 1, 1) diff --git a/providers/standard/tests/unit/standard/sensors/test_python.py b/providers/standard/tests/unit/standard/sensors/test_python.py index c363af3a9c4c7..4d3ede65f89ec 100644 --- a/providers/standard/tests/unit/standard/sensors/test_python.py +++ b/providers/standard/tests/unit/standard/sensors/test_python.py @@ -24,8 +24,8 @@ from airflow.exceptions import AirflowSensorTimeout from airflow.providers.standard.sensors.python import PythonSensor -from airflow.sensors.base import PokeReturnValue +from tests_common.test_utils.version_compat import PokeReturnValue from unit.standard.operators.test_python import BasePythonTest pytestmark = pytest.mark.db_test diff --git a/providers/standard/tests/unit/standard/sensors/test_time_delta.py b/providers/standard/tests/unit/standard/sensors/test_time_delta.py index 3ba3d14d0dec5..8e85be40f9c3f 100644 --- a/providers/standard/tests/unit/standard/sensors/test_time_delta.py +++ b/providers/standard/tests/unit/standard/sensors/test_time_delta.py @@ -19,7 +19,6 @@ from datetime import timedelta from typing import Any -from unittest import mock import pendulum import pytest @@ -34,16 +33,14 @@ WaitSensor, ) from airflow.providers.standard.triggers.temporal import DateTimeTrigger -from airflow.utils import timezone -from airflow.utils.timezone import datetime from airflow.utils.types import DagRunType from tests_common.test_utils import db -from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS +from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS, timezone pytestmark = pytest.mark.db_test -DEFAULT_DATE = datetime(2015, 1, 1) +DEFAULT_DATE = timezone.datetime(2015, 1, 1) DEV_NULL = "/dev/null" TEST_DAG_ID = "unit_tests" @@ -64,13 +61,13 @@ def setup_method(self): self.dagbag = DagBag(dag_folder=DEV_NULL, include_examples=False) self.dag = DAG(TEST_DAG_ID, schedule=timedelta(days=1), start_date=DEFAULT_DATE) - def test_timedelta_sensor(self): + def test_timedelta_sensor(self, mocker): op = TimeDeltaSensor(task_id="timedelta_sensor_check", delta=timedelta(seconds=2), dag=self.dag) - op.execute({"dag_run": mock.MagicMock(run_after=DEFAULT_DATE), "data_interval_end": DEFAULT_DATE}) + op.execute({"dag_run": mocker.MagicMock(run_after=DEFAULT_DATE), "data_interval_end": DEFAULT_DATE}) @pytest.mark.parametrize( - "run_after, interval_end", + ("run_after", "interval_end"), [ (timezone.utcnow() + timedelta(days=1), timezone.utcnow() + timedelta(days=2)), (timezone.utcnow() + timedelta(days=1), None), @@ -108,7 +105,7 @@ def test_timedelta_sensor_run_after_vs_interval(run_after, interval_end, dag_mak @pytest.mark.parametrize( - "run_after, interval_end", + ("run_after", "interval_end"), [ (timezone.utcnow() + timedelta(days=1), timezone.utcnow() + timedelta(days=2)), (timezone.utcnow() + timedelta(days=1), None), @@ -168,8 +165,8 @@ def setup_method(self): "should_defer", [False, True], ) - @mock.patch(DEFER_PATH) - def test_timedelta_sensor(self, defer_mock, should_defer): + def test_timedelta_sensor(self, mocker, should_defer): + defer_mock = mocker.patch(DEFER_PATH) delta = timedelta(hours=1) with pytest.warns(AirflowProviderDeprecationWarning): op = TimeDeltaSensorAsync(task_id="timedelta_sensor_check", delta=delta, dag=self.dag) @@ -187,9 +184,9 @@ def test_timedelta_sensor(self, defer_mock, should_defer): "should_defer", [False, True], ) - @mock.patch(DEFER_PATH) - @mock.patch("airflow.providers.standard.sensors.time_delta.sleep") - def test_wait_sensor(self, sleep_mock, defer_mock, should_defer): + def test_wait_sensor(self, mocker, should_defer): + defer_mock = mocker.patch(DEFER_PATH) + sleep_mock = mocker.patch("airflow.providers.standard.sensors.time_delta.sleep") wait_time = timedelta(seconds=30) op = WaitSensor( task_id="wait_sensor_check", time_to_wait=wait_time, dag=self.dag, deferrable=should_defer @@ -203,7 +200,7 @@ def test_wait_sensor(self, sleep_mock, defer_mock, should_defer): sleep_mock.assert_called_once_with(30) @pytest.mark.parametrize( - "run_after, interval_end", + ("run_after", "interval_end"), [ (timezone.utcnow() + timedelta(days=1), timezone.utcnow() + timedelta(days=2)), (timezone.utcnow() + timedelta(days=1), None), diff --git a/providers/standard/tests/unit/standard/sensors/test_weekday.py b/providers/standard/tests/unit/standard/sensors/test_weekday.py index b86349a032e93..3f7fa41ea26e1 100644 --- a/providers/standard/tests/unit/standard/sensors/test_weekday.py +++ b/providers/standard/tests/unit/standard/sensors/test_weekday.py @@ -26,18 +26,16 @@ from airflow.models.dag import DAG from airflow.providers.standard.sensors.weekday import DayOfWeekSensor from airflow.providers.standard.utils.weekday import WeekDay -from airflow.utils import timezone -from airflow.utils.timezone import datetime from tests_common.test_utils import db -from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS +from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS, timezone pytestmark = pytest.mark.db_test -DEFAULT_DATE = datetime(2018, 12, 10) -WEEKDAY_DATE = datetime(2018, 12, 20) -WEEKEND_DATE = datetime(2018, 12, 22) +DEFAULT_DATE = timezone.datetime(2018, 12, 10) +WEEKDAY_DATE = timezone.datetime(2018, 12, 20) +WEEKEND_DATE = timezone.datetime(2018, 12, 22) TEST_DAG_ID = "weekday_sensor_dag" DEV_NULL = "/dev/null" TEST_CASE_WEEKDAY_SENSOR_TRUE = { diff --git a/pyproject.toml b/pyproject.toml index d80c449f782b2..dbec1e042423e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -754,6 +754,7 @@ testing = ["dev", "providers.tests", "tests_common", "tests", "system", "unit", "kubernetes-tests/*" = ["D", "TID253", "S101", "TRY002"] "helm-tests/*" = ["D", "TID253", "S101", "TRY002"] "providers/**/tests/*" = ["D", "TID253", "S101", "TRY002"] +"performance/tests/*" = ["S101"] # All of the modules which have an extra license header (i.e. that we copy from another project) need to # ignore E402 -- module level import not at top level