Ejemplo n.º 1
0
 def test_job_delete_should_throw_ex_when_name_none(self, mock_hook):
     with self.assertRaises(AirflowException) as cm:
         op = CloudDataTransferServiceDeleteJobOperator(job_name="", task_id='task-id')
         op.execute(None)
     err = cm.exception
     self.assertIn("The required parameter 'job_name' is empty or None", str(err))
     mock_hook.assert_not_called()
 def test_job_delete(self, mock_hook):
     op = CloudDataTransferServiceDeleteJobOperator(
         job_name=JOB_NAME, project_id=GCP_PROJECT_ID, task_id='task-id')
     op.execute(None)
     mock_hook.assert_called_once_with(api_version='v1',
                                       gcp_conn_id='google_cloud_default')
     mock_hook.return_value.delete_transfer_job.assert_called_once_with(
         job_name=JOB_NAME, project_id=GCP_PROJECT_ID)
Ejemplo n.º 3
0
 def test_job_delete_with_templates(self, _):
     dag_id = 'test_dag_id'
     args = {'start_date': DEFAULT_DATE}
     self.dag = DAG(dag_id, default_args=args)  # pylint: disable=attribute-defined-outside-init
     op = CloudDataTransferServiceDeleteJobOperator(
         job_name='{{ dag.dag_id }}',
         gcp_conn_id='{{ dag.dag_id }}',
         api_version='{{ dag.dag_id }}',
         task_id=TASK_ID,
         dag=self.dag,
     )
     ti = TaskInstance(op, DEFAULT_DATE)
     ti.render_templates()
     assert dag_id == getattr(op, 'job_name')
     assert dag_id == getattr(op, 'gcp_conn_id')
     assert dag_id == getattr(op, 'api_version')
        job_name=
        "{{task_instance.xcom_pull('create_transfer_job_from_aws')['name']}}",
        project_id=GCP_PROJECT_ID,
        expected_statuses={GcpTransferOperationStatus.SUCCESS},
        poke_interval=WAIT_FOR_OPERATION_POKE_INTERVAL,
    )
    # [END howto_operator_gcp_transfer_wait_operation]

    # [START howto_operator_gcp_transfer_cancel_operation]
    cancel_operation = CloudDataTransferServiceCancelOperationOperator(
        task_id="cancel_operation",
        operation_name="{{task_instance.xcom_pull("
        "'wait_for_second_operation_to_start', key='sensed_operations')[0]['name']}}",
    )
    # [END howto_operator_gcp_transfer_cancel_operation]

    # [START howto_operator_gcp_transfer_delete_job]
    delete_transfer_from_aws_job = CloudDataTransferServiceDeleteJobOperator(
        task_id="delete_transfer_from_aws_job",
        job_name=
        "{{task_instance.xcom_pull('create_transfer_job_from_aws')['name']}}",
        project_id=GCP_PROJECT_ID,
    )
    # [END howto_operator_gcp_transfer_delete_job]

    # fmt: off
    create_transfer_job_from_aws >> wait_for_operation_to_start >> pause_operation
    pause_operation >> list_operations >> get_operation >> resume_operation
    resume_operation >> wait_for_operation_to_end >> cancel_operation >> delete_transfer_from_aws_job
    # fmt: on