def test_execute(self, mock_temp, mock_hook_gcs, mock_hook_wasb):
op = AzureBlobStorageToGCSOperator(
wasb_conn_id=WASB_CONN_ID,
gcp_conn_id=GCP_CONN_ID,
blob_name=BLOB_NAME,
file_path=FILE_PATH,
container_name=CONTAINER_NAME,
bucket_name=BUCKET_NAME,
object_name=OBJECT_NAME,
filename=FILENAME,
gzip=GZIP,
delegate_to=DELEGATE_TO,
impersonation_chain=IMPERSONATION_CHAIN,
task_id=TASK_ID,
)
op.execute(context=None)
mock_hook_wasb.assert_called_once_with(wasb_conn_id=WASB_CONN_ID)
mock_hook_wasb.return_value.get_file.assert_called_once_with(
file_path=mock_temp.NamedTemporaryFile.return_value.__enter__.return_value.name,
container_name=CONTAINER_NAME,
blob_name=BLOB_NAME,
)
mock_hook_gcs.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID, delegate_to=DELEGATE_TO, impersonation_chain=IMPERSONATION_CHAIN
)
mock_hook_gcs.return_value.upload.assert_called_once_with(
bucket_name=BUCKET_NAME,
object_name=OBJECT_NAME,
gzip=GZIP,
filename=mock_temp.NamedTemporaryFile.return_value.__enter__.return_value.name,
)
def test_init(self):
operator = AzureBlobStorageToGCSOperator(
wasb_conn_id=WASB_CONN_ID,
gcp_conn_id=GCP_CONN_ID,
blob_name=BLOB_NAME,
file_path=FILE_PATH,
container_name=CONTAINER_NAME,
bucket_name=BUCKET_NAME,
object_name=OBJECT_NAME,
filename=FILENAME,
gzip=GZIP,
delegate_to=DELEGATE_TO,
impersonation_chain=IMPERSONATION_CHAIN,
task_id=TASK_ID,
)
self.assertEqual(operator.wasb_conn_id, WASB_CONN_ID)
self.assertEqual(operator.blob_name, BLOB_NAME)
self.assertEqual(operator.file_path, FILE_PATH)
self.assertEqual(operator.container_name, CONTAINER_NAME)
self.assertEqual(operator.gcp_conn_id, GCP_CONN_ID)
self.assertEqual(operator.bucket_name, BUCKET_NAME)
self.assertEqual(operator.object_name, OBJECT_NAME)
self.assertEqual(operator.filename, FILENAME)
self.assertEqual(operator.gzip, GZIP)
self.assertEqual(operator.delegate_to, DELEGATE_TO)
self.assertEqual(operator.impersonation_chain, IMPERSONATION_CHAIN)
self.assertEqual(operator.task_id, TASK_ID)
def test_init(self):
operator = AzureBlobStorageToGCSOperator(
wasb_conn_id=WASB_CONN_ID,
gcp_conn_id=GCP_CONN_ID,
blob_name=BLOB_NAME,
file_path=FILE_PATH,
container_name=CONTAINER_NAME,
bucket_name=BUCKET_NAME,
object_name=OBJECT_NAME,
filename=FILENAME,
gzip=GZIP,
delegate_to=DELEGATE_TO,
impersonation_chain=IMPERSONATION_CHAIN,
task_id=TASK_ID,
)
assert operator.wasb_conn_id == WASB_CONN_ID
assert operator.blob_name == BLOB_NAME
assert operator.file_path == FILE_PATH
assert operator.container_name == CONTAINER_NAME
assert operator.gcp_conn_id == GCP_CONN_ID
assert operator.bucket_name == BUCKET_NAME
assert operator.object_name == OBJECT_NAME
assert operator.filename == FILENAME
assert operator.gzip == GZIP
assert operator.delegate_to == DELEGATE_TO
assert operator.impersonation_chain == IMPERSONATION_CHAIN
assert operator.task_id == TASK_ID
# [START how_to_wait_for_blob]
wait_for_blob = WasbBlobSensor(
task_id="wait_for_blob",
wasb_conn_id="wasb_default",
container_name=AZURE_CONTAINER_NAME,
blob_name=BLOB_NAME,
)
# [END how_to_wait_for_blob]
# [START how_to_azure_blob_to_gcs]
transfer_files_to_gcs = AzureBlobStorageToGCSOperator(
task_id="transfer_files_to_gcs",
# AZURE args
wasb_conn_id="wasb_default",
container_name=AZURE_CONTAINER_NAME,
blob_name=BLOB_NAME,
file_path=GCP_OBJECT_NAME,
# GCP args
gcp_conn_id="google_cloud_default",
bucket_name=GCP_BUCKET_NAME,
object_name=GCP_OBJECT_NAME,
filename=GCP_BUCKET_FILE_PATH,
gzip=False,
delegate_to=None,
impersonation_chain=None,
)
# [END how_to_azure_blob_to_gcs]
wait_for_blob >> transfer_files_to_gcs
GCP_OBJECT_NAME = os.environ.get("GCP_OBJECT_NAME", "file.txt")
# [START how_to_azure_blob_to_gcs]
with DAG(
"example_azure_blob_to_gcs",
schedule_interval=None,
start_date=datetime(2021, 1, 1), # Override to match your needs
default_args={
"container_name": AZURE_CONTAINER_NAME,
"blob_name": BLOB_NAME
},
) as dag:
wait_for_blob = WasbBlobSensor(task_id="wait_for_blob")
transfer_files_to_gcs = AzureBlobStorageToGCSOperator(
task_id="transfer_files_to_gcs",
# AZURE arg
file_path=GCP_OBJECT_NAME,
# GCP args
bucket_name=GCP_BUCKET_NAME,
object_name=GCP_OBJECT_NAME,
filename=GCP_BUCKET_FILE_PATH,
gzip=False,
delegate_to=None,
impersonation_chain=None,
)
# [END how_to_azure_blob_to_gcs]
wait_for_blob >> transfer_files_to_gcs