Skip to content

Commit

Permalink
Deprecate get_hook in OSSKeySensor and use hook instead (#34426)
Browse files Browse the repository at this point in the history
* Deprecate get_hook in OSSKeySensor and use hook instead

* use AirflowProviderDeprecationWarning

* Fix unit tests
  • Loading branch information
hussein-awala committed Sep 18, 2023
1 parent 169ce92 commit f8ae8db
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 12 deletions.
19 changes: 11 additions & 8 deletions airflow/providers/alibaba/cloud/sensors/oss_key.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,9 @@
from typing import TYPE_CHECKING, Sequence
from urllib.parse import urlsplit

from airflow.exceptions import AirflowException, AirflowSkipException
from deprecated.classic import deprecated

from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning, AirflowSkipException
from airflow.providers.alibaba.cloud.hooks.oss import OSSHook
from airflow.sensors.base import BaseSensorOperator

Expand Down Expand Up @@ -59,7 +61,6 @@ def __init__(
self.bucket_key = bucket_key
self.region = region
self.oss_conn_id = oss_conn_id
self.hook: OSSHook | None = None

def poke(self, context: Context):
"""
Expand Down Expand Up @@ -92,13 +93,15 @@ def poke(self, context: Context):
raise AirflowException(message)

self.log.info("Poking for key : oss://%s/%s", self.bucket_name, self.bucket_key)
return self.get_hook.object_exists(key=self.bucket_key, bucket_name=self.bucket_name)
return self.hook.object_exists(key=self.bucket_key, bucket_name=self.bucket_name)

@cached_property
@property
@deprecated(reason="use `hook` property instead.", category=AirflowProviderDeprecationWarning)
def get_hook(self) -> OSSHook:
"""Create and return an OSSHook."""
if self.hook:
return self.hook

self.hook = OSSHook(oss_conn_id=self.oss_conn_id, region=self.region)
return self.hook

@cached_property
def hook(self) -> OSSHook:
"""Create and return an OSSHook."""
return OSSHook(oss_conn_id=self.oss_conn_id, region=self.region)
8 changes: 4 additions & 4 deletions tests/providers/alibaba/cloud/sensors/test_oss_key.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,10 +50,10 @@ def oss_key_sensor():
class TestOSSKeySensor:
@mock.patch(f"{MODULE_NAME}.OSSHook")
def test_get_hook(self, mock_service, oss_key_sensor):
oss_key_sensor.get_hook()
oss_key_sensor.hook
mock_service.assert_called_once_with(oss_conn_id=MOCK_OSS_CONN_ID, region=MOCK_REGION)

@mock.patch(f"{MODULE_NAME}.OSSKeySensor.get_hook", new_callable=PropertyMock)
@mock.patch(f"{MODULE_NAME}.OSSKeySensor.hook", new_callable=PropertyMock)
def test_poke_exsiting_key(self, mock_service, oss_key_sensor):
# Given
mock_service.return_value.object_exists.return_value = True
Expand All @@ -65,7 +65,7 @@ def test_poke_exsiting_key(self, mock_service, oss_key_sensor):
assert res is True
mock_service.return_value.object_exists.assert_called_once_with(key=MOCK_KEY, bucket_name=MOCK_BUCKET)

@mock.patch(f"{MODULE_NAME}.OSSKeySensor.get_hook", new_callable=PropertyMock)
@mock.patch(f"{MODULE_NAME}.OSSKeySensor.hook", new_callable=PropertyMock)
def test_poke_non_exsiting_key(self, mock_service, oss_key_sensor):
# Given
mock_service.return_value.object_exists.return_value = False
Expand All @@ -80,7 +80,7 @@ def test_poke_non_exsiting_key(self, mock_service, oss_key_sensor):
@pytest.mark.parametrize(
"soft_fail, expected_exception", ((False, AirflowException), (True, AirflowSkipException))
)
@mock.patch(f"{MODULE_NAME}.OSSKeySensor.get_hook", new_callable=PropertyMock)
@mock.patch(f"{MODULE_NAME}.OSSKeySensor.hook", new_callable=PropertyMock)
def test_poke_without_bucket_name(
self, mock_service, oss_key_sensor, soft_fail: bool, expected_exception: AirflowException
):
Expand Down

0 comments on commit f8ae8db

Please sign in to comment.