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', })