def test_poke(self, mock_hook):
mock_instance = mock_hook.return_value
sensor = WasbBlobSensor(
task_id='wasb_sensor', dag=self.dag, check_options={'timeout': 2}, **self._config
)
sensor.poke(None)
mock_instance.check_for_blob.assert_called_once_with('container', 'blob', timeout=2)
def test_init(self):
sensor = WasbBlobSensor(task_id='wasb_sensor_1', dag=self.dag, **self._config)
self.assertEqual(sensor.container_name, self._config['container_name'])
self.assertEqual(sensor.blob_name, self._config['blob_name'])
self.assertEqual(sensor.wasb_conn_id, self._config['wasb_conn_id'])
self.assertEqual(sensor.check_options, {})
self.assertEqual(sensor.timeout, self._config['timeout'])
sensor = WasbBlobSensor(
task_id='wasb_sensor_2', dag=self.dag, check_options={'timeout': 2}, **self._config
)
self.assertEqual(sensor.check_options, {'timeout': 2})
def test_init(self):
sensor = WasbBlobSensor(task_id='wasb_sensor_1',
dag=self.dag,
**self._config)
assert sensor.container_name == self._config['container_name']
assert sensor.blob_name == self._config['blob_name']
assert sensor.wasb_conn_id == self._config['wasb_conn_id']
assert sensor.check_options == {}
assert sensor.timeout == self._config['timeout']
sensor = WasbBlobSensor(task_id='wasb_sensor_2',
dag=self.dag,
check_options={'timeout': 2},
**self._config)
assert sensor.check_options == {'timeout': 2}
BLOB_NAME = os.environ.get("AZURE_BLOB_NAME", "file.txt")
AZURE_CONTAINER_NAME = os.environ.get("AZURE_CONTAINER_NAME", "airflow")
GCP_BUCKET_FILE_PATH = os.environ.get("GCP_BUCKET_FILE_PATH", "file.txt")
GCP_BUCKET_NAME = os.environ.get("GCP_BUCKET_NAME", "azure_bucket")
GCP_OBJECT_NAME = os.environ.get("GCP_OBJECT_NAME", "file.txt")
with DAG(
"example_azure_blob_to_gcs",
schedule_interval=None,
start_date=days_ago(1), # Override to match your needs
) as dag:
# [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,
GCP_BUCKET_FILE_PATH = os.environ.get("GCP_BUCKET_FILE_PATH", "file.txt")
GCP_BUCKET_NAME = os.environ.get("GCP_BUCKET_NAME", "INVALID BUCKET NAME")
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]