Skip to content

Commit

Permalink
Fix GCS sensor system tests failing with DebugExecutor (#26742)
Browse files Browse the repository at this point in the history
  • Loading branch information
bhirsz authored Sep 28, 2022
1 parent 4521188 commit dce2755
Show file tree
Hide file tree
Showing 3 changed files with 24 additions and 4 deletions.
2 changes: 1 addition & 1 deletion airflow/providers/google/cloud/sensors/gcs.py
Original file line number Diff line number Diff line change
Expand Up @@ -241,7 +241,7 @@ class GCSUploadSessionCompleteSensor(BaseSensorOperator):
"""
Checks for changes in the number of objects at prefix in Google Cloud Storage
bucket and returns True if the inactivity period has passed with no
increase in the number of objects. Note, this sensor will no behave correctly
increase in the number of objects. Note, this sensor will not behave correctly
in reschedule mode, as the state of the listed objects in the GCS bucket will
be lost between rescheduled invocations.
Expand Down
2 changes: 1 addition & 1 deletion airflow/sensors/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,7 @@ class BaseSensorOperator(BaseOperator, SkipMixin):
:param soft_fail: Set to true to mark the task as SKIPPED on failure
:param poke_interval: Time in seconds that the job should wait in
between each tries
between each try
:param timeout: Time, in seconds before the task times out and fails.
:param mode: How the sensor operates.
Options are: ``{ poke | reschedule }``, default is ``poke``.
Expand Down
24 changes: 22 additions & 2 deletions tests/system/providers/google/cloud/gcs/example_gcs_sensor.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,26 @@
UPLOAD_FILE_PATH = str(Path(__file__).parent / "resources" / FILE_NAME)


def workaround_in_debug_executor(cls):
"""
DebugExecutor change sensor mode from poke to reschedule. Some sensors don't work correctly
in reschedule mode. They are decorated with `poke_mode_only` decorator to fail when mode is changed.
This method creates dummy property to overwrite it and force poke method to always return True.
"""
cls.mode = dummy_mode_property()
cls.poke = lambda self, ctx: True


def dummy_mode_property():
def mode_getter(self):
return self._mode

def mode_setter(self, value):
self._mode = value

return property(mode_getter, mode_setter)


with models.DAG(
DAG_ID,
schedule='@once',
Expand All @@ -58,6 +78,8 @@
task_id="create_bucket", bucket_name=BUCKET_NAME, project_id=PROJECT_ID
)

workaround_in_debug_executor(GCSUploadSessionCompleteSensor)

# [START howto_sensor_gcs_upload_session_complete_task]
gcs_upload_session_complete = GCSUploadSessionCompleteSensor(
bucket=BUCKET_NAME,
Expand Down Expand Up @@ -89,7 +111,6 @@
gcs_object_exists = GCSObjectExistenceSensor(
bucket=BUCKET_NAME,
object=FILE_NAME,
mode='poke',
task_id="gcs_object_exists_task",
)
# [END howto_sensor_object_exists_task]
Expand All @@ -98,7 +119,6 @@
gcs_object_with_prefix_exists = GCSObjectsWithPrefixExistenceSensor(
bucket=BUCKET_NAME,
prefix=FILE_NAME[:5],
mode='poke',
task_id="gcs_object_with_prefix_exists_task",
)
# [END howto_sensor_object_with_prefix_exists_task]
Expand Down

0 comments on commit dce2755

Please sign in to comment.