def test_wait_for_status_after_retry(self, mock_tool):
operations_set = [
[{'metadata': {'status': GcpTransferOperationStatus.SUCCESS}}],
[{'metadata': {'status': GcpTransferOperationStatus.SUCCESS}}],
]
mock_tool.return_value.list_transfer_operations.side_effect = operations_set
mock_tool.operations_contain_expected_statuses.side_effect = [False, True]
op = CloudDataTransferServiceJobStatusSensor(
task_id='task-id',
job_name='job-name',
project_id='project-id',
expected_statuses=GcpTransferOperationStatus.SUCCESS,
)
context = {'ti': (mock.Mock(**{'xcom_push.return_value': None}))}
result = op.poke(context)
self.assertFalse(result)
mock_tool.operations_contain_expected_statuses.assert_called_once_with(
operations=operations_set[0], expected_statuses={GcpTransferOperationStatus.SUCCESS}
)
mock_tool.operations_contain_expected_statuses.reset_mock()
result = op.poke(context)
self.assertTrue(result)
mock_tool.operations_contain_expected_statuses.assert_called_once_with(
operations=operations_set[1], expected_statuses={GcpTransferOperationStatus.SUCCESS}
)
def test_wait_for_status_success(self, mock_tool):
operations = [{
'metadata': {
'status': GcpTransferOperationStatus.SUCCESS
}
}]
mock_tool.return_value.list_transfer_operations.return_value = operations
mock_tool.operations_contain_expected_statuses.return_value = True
op = CloudDataTransferServiceJobStatusSensor(
task_id='task-id',
job_name='job-name',
project_id='project-id',
expected_statuses=GcpTransferOperationStatus.SUCCESS,
)
context = {'ti': (mock.Mock(**{'xcom_push.return_value': None}))}
result = op.poke(context)
mock_tool.return_value.list_transfer_operations.assert_called_once_with(
request_filter={
'project_id': 'project-id',
'job_names': ['job-name']
})
mock_tool.operations_contain_expected_statuses.assert_called_once_with(
operations=operations,
expected_statuses={GcpTransferOperationStatus.SUCCESS})
self.assertTrue(result)
def test_wait_for_status_normalize_status(self, expected_status,
received_status, mock_tool):
operations = [{
'metadata': {
'status': GcpTransferOperationStatus.SUCCESS
}
}]
mock_tool.return_value.list_transfer_operations.return_value = operations
mock_tool.operations_contain_expected_statuses.side_effect = [
False, True
]
op = CloudDataTransferServiceJobStatusSensor(
task_id='task-id',
job_name='job-name',
project_id='project-id',
expected_statuses=expected_status,
)
context = {'ti': (mock.Mock(**{'xcom_push.return_value': None}))}
result = op.poke(context)
self.assertFalse(result)
mock_tool.operations_contain_expected_statuses.assert_called_once_with(
operations=operations, expected_statuses=received_status)
def test_wait_for_status_success_default_expected_status(self, mock_tool):
op = CloudDataTransferServiceJobStatusSensor(
task_id='task-id',
job_name='job-name',
project_id='project-id',
expected_statuses=GcpTransferOperationStatus.SUCCESS,
)
context = {'ti': (mock.Mock(**{'xcom_push.return_value': None}))}
result = op.poke(context)
mock_tool.operations_contain_expected_statuses.assert_called_once_with(
operations=mock.ANY, expected_statuses={GcpTransferOperationStatus.SUCCESS}
)
self.assertTrue(result)
with models.DAG(
'example_gcp_transfer_aws',
schedule_interval=None, # Override to match your needs
start_date=days_ago(1),
tags=['example'],
) as dag:
# [START howto_operator_gcp_transfer_create_job]
create_transfer_job_from_aws = CloudDataTransferServiceCreateJobOperator(
task_id="create_transfer_job_from_aws", body=aws_to_gcs_transfer_body)
# [END howto_operator_gcp_transfer_create_job]
wait_for_operation_to_start = CloudDataTransferServiceJobStatusSensor(
task_id="wait_for_operation_to_start",
job_name=
"{{task_instance.xcom_pull('create_transfer_job_from_aws')['name']}}",
project_id=GCP_PROJECT_ID,
expected_statuses={GcpTransferOperationStatus.IN_PROGRESS},
poke_interval=WAIT_FOR_OPERATION_POKE_INTERVAL,
)
# [START howto_operator_gcp_transfer_pause_operation]
pause_operation = CloudDataTransferServicePauseOperationOperator(
task_id="pause_operation",
operation_name=
"{{task_instance.xcom_pull('wait_for_operation_to_start', "
"key='sensed_operations')[0]['name']}}",
)
# [END howto_operator_gcp_transfer_pause_operation]
# [START howto_operator_gcp_transfer_list_operations]
list_operations = CloudDataTransferServiceListOperationsOperator(
) as dag:
create_transfer = CloudDataTransferServiceCreateJobOperator(
task_id="create_transfer", body=gcs_to_gcs_transfer_body)
# [START howto_operator_gcp_transfer_update_job]
update_transfer = CloudDataTransferServiceUpdateJobOperator(
task_id="update_transfer",
job_name="{{task_instance.xcom_pull('create_transfer')['name']}}",
body=update_body,
)
# [END howto_operator_gcp_transfer_update_job]
wait_for_transfer = CloudDataTransferServiceJobStatusSensor(
task_id="wait_for_transfer",
job_name="{{task_instance.xcom_pull('create_transfer')['name']}}",
project_id=GCP_PROJECT_ID,
expected_statuses={GcpTransferOperationStatus.SUCCESS},
)
list_operations = CloudDataTransferServiceListOperationsOperator(
task_id="list_operations",
request_filter={
FILTER_PROJECT_ID:
GCP_PROJECT_ID,
FILTER_JOB_NAMES:
["{{task_instance.xcom_pull('create_transfer')['name']}}"],
},
)
get_operation = CloudDataTransferServiceGetOperationOperator(
task_id="get_operation",