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",