diff --git a/providers/http/src/airflow/providers/http/operators/http.py b/providers/http/src/airflow/providers/http/operators/http.py index ca36ac471c7fd..6e6c48ea94a90 100644 --- a/providers/http/src/airflow/providers/http/operators/http.py +++ b/providers/http/src/airflow/providers/http/operators/http.py @@ -28,8 +28,8 @@ from airflow.configuration import conf from airflow.exceptions import AirflowException from airflow.hooks.base import BaseHook -from airflow.models import BaseOperator from airflow.providers.http.triggers.http import HttpTrigger, serialize_auth_type +from airflow.providers.http.version_compat import BaseOperator from airflow.utils.helpers import merge_dicts if TYPE_CHECKING: diff --git a/providers/http/src/airflow/providers/http/sensors/http.py b/providers/http/src/airflow/providers/http/sensors/http.py index a9d9274afd109..f395ab2bb8ca3 100644 --- a/providers/http/src/airflow/providers/http/sensors/http.py +++ b/providers/http/src/airflow/providers/http/sensors/http.py @@ -25,12 +25,7 @@ from airflow.exceptions import AirflowException from airflow.providers.http.hooks.http import HttpHook from airflow.providers.http.triggers.http import HttpSensorTrigger -from airflow.providers.http.version_compat import AIRFLOW_V_3_0_PLUS - -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk import BaseSensorOperator -else: - from airflow.sensors.base import BaseSensorOperator # type: ignore[no-redef] +from airflow.providers.http.version_compat import AIRFLOW_V_3_0_PLUS, BaseSensorOperator if TYPE_CHECKING: try: diff --git a/providers/http/src/airflow/providers/http/version_compat.py b/providers/http/src/airflow/providers/http/version_compat.py index 48d122b669696..fde11c60d3655 100644 --- a/providers/http/src/airflow/providers/http/version_compat.py +++ b/providers/http/src/airflow/providers/http/version_compat.py @@ -33,3 +33,11 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3, 0, 0) + +if AIRFLOW_V_3_0_PLUS: + from airflow.sdk import BaseOperator, BaseSensorOperator +else: + from airflow.models import BaseOperator # type: ignore[no-redef] + from airflow.sensors.base import BaseSensorOperator # type: ignore[no-redef] + +__all__ = ["AIRFLOW_V_3_0_PLUS", "BaseOperator", "BaseSensorOperator"] diff --git a/providers/http/tests/unit/http/sensors/test_http.py b/providers/http/tests/unit/http/sensors/test_http.py index b2b74c1dc5bfa..1f3c0a32609d7 100644 --- a/providers/http/tests/unit/http/sensors/test_http.py +++ b/providers/http/tests/unit/http/sensors/test_http.py @@ -286,11 +286,6 @@ def mount(self, prefix, adapter): class TestHttpOpSensor: - def setup_method(self): - args = {"owner": "airflow", "start_date": DEFAULT_DATE_ISO} - dag = DAG(TEST_DAG_ID, schedule=None, default_args=args) - self.dag = dag - @mock.patch("airflow.providers.http.hooks.http.Session", FakeSession) def test_get(self): op = HttpOperator( @@ -299,9 +294,8 @@ def test_get(self): endpoint="/search", data={"client": "ubuntu", "q": "airflow"}, headers={}, - dag=self.dag, ) - op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True) + op.execute({}) @mock.patch("airflow.providers.http.hooks.http.Session", FakeSession) def test_get_response_check(self): @@ -312,9 +306,8 @@ def test_get_response_check(self): data={"client": "ubuntu", "q": "airflow"}, response_check=lambda response: ("apache/airflow" in response.text), headers={}, - dag=self.dag, ) - op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True) + op.execute({}) @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Test only for Airflow 3.0+") @mock.patch("airflow.providers.http.hooks.http.Session", FakeSession) @@ -328,13 +321,13 @@ def test_sensor(self, run_task): response_check=lambda response: f"apache/airflow/{DEFAULT_DATE:%Y-%m-%d}" in response.text, poke_interval=5, timeout=15, - dag=self.dag, ) run_task(sensor) @pytest.mark.skipif(AIRFLOW_V_3_0_PLUS, reason="Test only for Airflow < 3.0") @mock.patch("airflow.providers.http.hooks.http.Session", FakeSession) def test_sensor_af2(self): + dag = DAG(TEST_DAG_ID, schedule=None) sensor = HttpSensor( task_id="http_sensor_check", http_conn_id="http_default", @@ -344,7 +337,7 @@ def test_sensor_af2(self): response_check=lambda response: f"apache/airflow/{DEFAULT_DATE:%Y-%m-%d}" in response.text, poke_interval=5, timeout=15, - dag=self.dag, + dag=dag, ) sensor.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)