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_subprocess.PIPE = "pipe"
mock_subprocess.STDOUT = "stdout"
mock_subprocess.Popen.return_value.stdout.readline = lambda: b""
mock_subprocess.Popen.return_value.wait.return_value = None
mock_subprocess.Popen.return_value.returncode = 0
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,