def test_execute(self, mock_hook, mock_subprocess, mock_tempfile):
source_bucket = TEST_BUCKET
source_object = "test.txt"
destination_bucket = TEST_BUCKET + "-dest"
destination_object = "transformed_test.txt"
transform_script = "script.py"
source = "source"
destination = "destination"
# Mock the name attribute...
mock1 = mock.Mock()
mock2 = mock.Mock()
mock1.name = source
mock2.name = destination
mock_tempfile.return_value.__enter__.side_effect = [mock1, mock2]
mock_proc = mock.MagicMock()
mock_proc.returncode = 0
mock_proc.stdout.readline = lambda: b""
mock_proc.wait.return_value = None
mock_popen = mock.MagicMock()
mock_popen.return_value.__enter__.return_value = mock_proc
mock_subprocess.Popen = mock_popen
mock_subprocess.PIPE = "pipe"
mock_subprocess.STDOUT = "stdout"
op = GCSFileTransformOperator(
task_id=TASK_ID,
source_bucket=source_bucket,
source_object=source_object,
destination_object=destination_object,
destination_bucket=destination_bucket,
transform_script=transform_script,
)
op.execute(None)
mock_hook.return_value.download.assert_called_once_with(
bucket_name=source_bucket,
object_name=source_object,
filename=source)
mock_subprocess.Popen.assert_called_once_with(
args=[transform_script, source, destination],
stdout="pipe",
stderr="stdout",
close_fds=True,
)
mock_hook.return_value.upload.assert_called_with(
bucket_name=destination_bucket,
object_name=destination_object,
filename=destination,
)
list_buckets_result = BashOperator(
task_id="list_buckets_result",
bash_command="echo \"{{ task_instance.xcom_pull('list_buckets') }}\"",
)
upload_file = LocalFilesystemToGCSOperator(
task_id="upload_file",
src=PATH_TO_UPLOAD_FILE,
dst=BUCKET_FILE_LOCATION,
bucket=BUCKET_1,
)
transform_file = GCSFileTransformOperator(
task_id="transform_file",
source_bucket=BUCKET_1,
source_object=BUCKET_FILE_LOCATION,
transform_script=["python", PATH_TO_TRANSFORM_SCRIPT]
)
# [START howto_operator_gcs_bucket_create_acl_entry_task]
gcs_bucket_create_acl_entry_task = GCSBucketCreateAclEntryOperator(
bucket=BUCKET_1,
entity=GCS_ACL_ENTITY,
role=GCS_ACL_BUCKET_ROLE,
task_id="gcs_bucket_create_acl_entry_task",
)
# [END howto_operator_gcs_bucket_create_acl_entry_task]
# [START howto_operator_gcs_object_create_acl_entry_task]
gcs_object_create_acl_entry_task = GCSObjectCreateAclEntryOperator(
bucket=BUCKET_1,
object_name=BUCKET_FILE_LOCATION,
with DAG('demo_dag', schedule_interval=None, default_args=default_args) as dag:
wait_task = GoogleCloudStoragePrefixSensor(
task_id='filesensor',
bucket='{{var.value.gcs_bucket}}',
prefix='{{var.value.gcs_file}}',
google_cloud_conn_id='google_cloud_default',
dag=dag)
transform_file = GCSFileTransformOperator(
task_id="transform_file",
source_bucket='{{var.value.gcs_bucket}}',
source_object='devfest',
destination_bucket='{{var.value.gcs_bucket}}',
destination_object='new_file',
transform_script=["cp"],
google_cloud_conn_id='google_cloud_default',
)
wait_task >> transform_file
for i in 1, 2:
transform_file >> PythonOperator(task_id='hello_world_' + str(i),
provide_context=True,
python_callable=demo,
op_kwargs={
"argument": 'hello world',
})