Migrate GCS sensor examples

Bartlomiej Hirsz · Bartlomiej Hirsz · commit 5bb09f6067d8 · 2022-04-12T08:58:13.000Z
Change-Id: Iec9be2369a93f2b9bab62f6e5eef298492852c90
diff --git a/tests/always/test_project_structure.py b/tests/always/test_project_structure.py
@@ -214,8 +214,6 @@ class TestGoogleProviderProjectStructure(unittest.TestCase):
         'airflow.providers.google.cloud.operators.dlp.CloudDLPDeleteDeidentifyTemplateOperator',
         'airflow.providers.google.cloud.operators.dlp.CloudDLPListDLPJobsOperator',
         'airflow.providers.google.cloud.operators.dlp.CloudDLPRedactImageOperator',
-        'airflow.providers.google.cloud.sensors.gcs.GCSObjectUpdateSensor',
-        'airflow.providers.google.cloud.sensors.gcs.GCSUploadSessionCompleteSensor',
     }
 
     def test_missing_example_for_operator(self):
diff --git a/tests/system/providers/google/gcs/example_gcs_sensor.py b/tests/system/providers/google/gcs/example_gcs_sensor.py
@@ -16,8 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """
-Example Airflow DAG for Google Cloud Storage GCSObjectExistenceSensor and
-GCSObjectsWithPrefixExistenceSensor sensors.
+Example Airflow DAG for Google Cloud Storage sensors.
 """
 
 import os
@@ -26,10 +25,13 @@
 
 from airflow import models
 from airflow.models.baseoperator import chain
+from airflow.operators.bash import BashOperator
 from airflow.providers.google.cloud.operators.gcs import GCSCreateBucketOperator, GCSDeleteBucketOperator
 from airflow.providers.google.cloud.sensors.gcs import (
     GCSObjectExistenceSensor,
     GCSObjectsWithPrefixExistenceSensor,
+    GCSObjectUpdateSensor,
+    GCSUploadSessionCompleteSensor,
 )
 from airflow.providers.google.cloud.transfers.local_to_gcs import LocalFilesystemToGCSOperator
 from airflow.utils.trigger_rule import TriggerRule
@@ -55,6 +57,26 @@
         task_id="create_bucket", bucket_name=BUCKET_NAME, project_id=PROJECT_ID
     )
 
+    # [START howto_sensor_gcs_upload_session_complete_task]
+    gcs_upload_session_complete = GCSUploadSessionCompleteSensor(
+        bucket=BUCKET_NAME,
+        prefix=FILE_NAME,
+        inactivity_period=15,
+        min_objects=1,
+        allow_delete=True,
+        previous_objects=set(),
+        task_id="gcs_upload_session_complete_task",
+    )
+    # [END howto_sensor_gcs_upload_session_complete_task]
+
+    # [START howto_sensor_object_update_exists_task]
+    gcs_update_object_exists = GCSObjectUpdateSensor(
+        bucket=BUCKET_NAME,
+        object=FILE_NAME,
+        task_id="gcs_object_update_sensor_task",
+    )
+    # [END howto_sensor_object_update_exists_task]
+
     upload_file = LocalFilesystemToGCSOperator(
         task_id="upload_file",
         src=UPLOAD_FILE_PATH,
@@ -84,15 +106,25 @@
         task_id="delete_bucket", bucket_name=BUCKET_NAME, trigger_rule=TriggerRule.ALL_DONE
     )
 
+    sleep = BashOperator(task_id='sleep', bash_command='sleep 5')
+
     chain(
         # TEST SETUP
         create_bucket,
+        sleep,
         upload_file,
         # TEST BODY
         [gcs_object_exists, gcs_object_with_prefix_exists],
         # TEST TEARDOWN
         delete_bucket,
     )
+    chain(
+        create_bucket,
+        # TEST BODY
+        gcs_upload_session_complete,
+        gcs_update_object_exists,
+        delete_bucket,
+    )
 
     from tests.system.utils.watcher import watcher
 

Original file line number	Diff line number	Diff line change
`@@ -214,8 +214,6 @@ class TestGoogleProviderProjectStructure(unittest.TestCase):`
`214`	`214`	`'airflow.providers.google.cloud.operators.dlp.CloudDLPDeleteDeidentifyTemplateOperator',`
`215`	`215`	`'airflow.providers.google.cloud.operators.dlp.CloudDLPListDLPJobsOperator',`
`216`	`216`	`'airflow.providers.google.cloud.operators.dlp.CloudDLPRedactImageOperator',`
`217`		`- 'airflow.providers.google.cloud.sensors.gcs.GCSObjectUpdateSensor',`
`218`		`- 'airflow.providers.google.cloud.sensors.gcs.GCSUploadSessionCompleteSensor',`
`219`	`217`	`}`
`220`	`218`
`221`	`219`	`def test_missing_example_for_operator(self):`