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 = GCPTransferServiceWaitForJobStatusSensor(
task_id='task-id',
operation_name='operation-name',
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_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_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 = GCPTransferServiceWaitForJobStatusSensor(
task_id='task-id',
operation_name='operation-name',
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_with(
filter={
'project_id': 'project-id',
'job_names': ['job-name']
})
mock_tool.operations_contain_expected_statuses.assert_called_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 = GCPTransferServiceWaitForJobStatusSensor(
task_id='task-id',
operation_name='operation-name',
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_with(
operations=operations, expected_statuses=received_status)
def test_wait_for_status_success_default_expected_status(self, mock_tool):
op = GCPTransferServiceWaitForJobStatusSensor(
task_id='task-id',
operation_name='operation-name',
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_with(
operations=mock.ANY, expected_statuses={GcpTransferOperationStatus.SUCCESS}
)
self.assertTrue(result)
with models.DAG(
'example_gcp_transfer',
default_args=default_args,
schedule_interval=None # Override to match your needs
) as dag:
# [START howto_operator_gcp_transfer_create_job]
create_transfer_job_from_aws = GcpTransferServiceJobCreateOperator(
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 = GCPTransferServiceWaitForJobStatusSensor(
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 = GcpTransferServiceOperationPauseOperator(
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_update_job]
update_job = GcpTransferServiceJobUpdateOperator(