Example #1
0
    def test_should_pass_argument_to_hook(self, mock_hook):
        task = GCSObjectUpdateSensor(
            task_id="task-id",
            bucket=TEST_BUCKET,
            object=TEST_OBJECT,
            google_cloud_conn_id=TEST_GCP_CONN_ID,
            delegate_to=TEST_DELEGATE_TO,
        )
        mock_hook.return_value.is_updated_after.return_value = True
        result = task.poke(mock.MagicMock())

        mock_hook.assert_called_once_with(
            delegate_to=TEST_DELEGATE_TO,
            google_cloud_storage_conn_id=TEST_GCP_CONN_ID
        )
        mock_hook.return_value.is_updated_after.assert_called_once_with(TEST_BUCKET, TEST_OBJECT, mock.ANY)
        self.assertEqual(True, result)
Example #2
0
    def test_should_pass_argument_to_hook(self, mock_hook):
        task = GCSObjectUpdateSensor(
            task_id="task-id",
            bucket=TEST_BUCKET,
            object=TEST_OBJECT,
            google_cloud_conn_id=TEST_GCP_CONN_ID,
            delegate_to=TEST_DELEGATE_TO,
            impersonation_chain=TEST_IMPERSONATION_CHAIN,
        )
        mock_hook.return_value.is_updated_after.return_value = True
        result = task.poke(mock.MagicMock())

        mock_hook.assert_called_once_with(
            delegate_to=TEST_DELEGATE_TO,
            gcp_conn_id=TEST_GCP_CONN_ID,
            impersonation_chain=TEST_IMPERSONATION_CHAIN,
        )
        mock_hook.return_value.is_updated_after.assert_called_once_with(
            TEST_BUCKET, TEST_OBJECT, mock.ANY)
        assert result is True
Example #3
0
    # [START howto_sensor_gcs_upload_session_complete_task]
    gcs_upload_session_complete = GCSUploadSessionCompleteSensor(
        bucket=BUCKET_1,
        prefix=PATH_TO_MANUAL_UPLOAD_FILE_PREFIX,
        inactivity_period=60,
        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_1,
        object=BUCKET_MANUAL_UPLOAD_FILE_LOCATION,
        task_id="gcs_object_update_sensor_task",
    )
    # [END howto_sensor_object_update_exists_task]

    delete_bucket = GCSDeleteBucketOperator(task_id="delete_bucket",
                                            bucket_name=BUCKET_1)

    create_bucket >> upload_file >> [
        gcs_object_exists, gcs_object_with_prefix_exists
    ] >> delete_bucket
    create_bucket >> gcs_upload_session_complete >> gcs_update_object_exists >> delete_bucket

if __name__ == '__main__':
    dag.clear()
    dag.run()