diff --git a/airflow/providers/alibaba/cloud/sensors/analyticdb_spark.py b/airflow/providers/alibaba/cloud/sensors/analyticdb_spark.py index fb6a962d43798..02bab417b6d7a 100644 --- a/airflow/providers/alibaba/cloud/sensors/analyticdb_spark.py +++ b/airflow/providers/alibaba/cloud/sensors/analyticdb_spark.py @@ -20,6 +20,9 @@ from functools import cached_property from typing import TYPE_CHECKING, Any, Sequence +from deprecated.classic import deprecated + +from airflow.exceptions import AirflowProviderDeprecationWarning from airflow.providers.alibaba.cloud.hooks.analyticdb_spark import AnalyticDBSparkHook, AppState from airflow.sensors.base import BaseSensorOperator @@ -50,19 +53,19 @@ def __init__( self.app_id = app_id self._region = region self._adb_spark_conn_id = adb_spark_conn_id - self._adb_spark_hook: AnalyticDBSparkHook | None = None @cached_property + def hook(self) -> AnalyticDBSparkHook: + """Get valid hook.""" + return AnalyticDBSparkHook(adb_spark_conn_id=self._adb_spark_conn_id, region=self._region) + + @deprecated(reason="use `hook` property instead.", category=AirflowProviderDeprecationWarning) def get_hook(self) -> AnalyticDBSparkHook: """Get valid hook.""" - if self._adb_spark_hook is None or not isinstance(self._adb_spark_hook, AnalyticDBSparkHook): - self._adb_spark_hook = AnalyticDBSparkHook( - adb_spark_conn_id=self._adb_spark_conn_id, region=self._region - ) - return self._adb_spark_hook + return self.hook def poke(self, context: Context) -> bool: app_id = self.app_id - state = self.get_hook.get_spark_state(app_id) + state = self.hook.get_spark_state(app_id) return AppState(state) in AnalyticDBSparkHook.TERMINAL_STATES diff --git a/tests/providers/alibaba/cloud/sensors/test_analyticdb_spark.py b/tests/providers/alibaba/cloud/sensors/test_analyticdb_spark.py index 8cef5175005f3..86c67190635ac 100644 --- a/tests/providers/alibaba/cloud/sensors/test_analyticdb_spark.py +++ b/tests/providers/alibaba/cloud/sensors/test_analyticdb_spark.py @@ -42,10 +42,10 @@ def setup_method(self): @mock.patch(ADB_SPARK_SENSOR_STRING.format("AnalyticDBSparkHook")) def test_get_hook(self, mock_service): """Test get_hook function works as expected.""" - self.sensor.get_hook() + self.sensor.hook mock_service.assert_called_once_with(adb_spark_conn_id=MOCK_ADB_SPARK_CONN_ID, region=MOCK_REGION) - @mock.patch(ADB_SPARK_SENSOR_STRING.format("AnalyticDBSparkSensor.get_hook")) + @mock.patch(ADB_SPARK_SENSOR_STRING.format("AnalyticDBSparkSensor.hook")) def test_poke_terminal_state(self, mock_service): """Test poke_terminal_state works as expected with COMPLETED application.""" # Given @@ -58,7 +58,7 @@ def test_poke_terminal_state(self, mock_service): assert res is True mock_service.get_spark_state.assert_called_once_with(MOCK_ADB_SPARK_ID) - @mock.patch(ADB_SPARK_SENSOR_STRING.format("AnalyticDBSparkSensor.get_hook")) + @mock.patch(ADB_SPARK_SENSOR_STRING.format("AnalyticDBSparkSensor.hook")) def test_poke_non_terminal_state(self, mock_service): """Test poke_terminal_state works as expected with RUNNING application.""" # Given