def test_operation_list(self, mock_hook): mock_hook.return_value.list_transfer_operations.return_value = [VALID_TRANSFER_JOB_GCS] op = GcpTransferServiceOperationsListOperator(request_filter=TEST_FILTER, task_id=TASK_ID) result = op.execute(None) mock_hook.assert_called_once_with(api_version='v1', gcp_conn_id='google_cloud_default') mock_hook.return_value.list_transfer_operations.assert_called_once_with(request_filter=TEST_FILTER) self.assertEqual(result, [VALID_TRANSFER_JOB_GCS])
def test_operation_list(self, mock_hook): mock_hook.return_value.list_transfer_operations.return_value = [VALID_TRANSFER_JOB_GCS] op = GcpTransferServiceOperationsListOperator(filter=TEST_FILTER, task_id=TASK_ID) result = op.execute(None) mock_hook.assert_called_once_with(api_version='v1', gcp_conn_id='google_cloud_default') mock_hook.return_value.list_transfer_operations.assert_called_once_with(filter=TEST_FILTER) self.assertEqual(result, [VALID_TRANSFER_JOB_GCS])
def test_templates(self, _): dag_id = 'test_dag_id' args = {'start_date': DEFAULT_DATE} self.dag = DAG(dag_id, default_args=args) op = GcpTransferServiceOperationsListOperator( filter={"job_names": ['{{ dag.dag_id }}']}, gcp_conn_id='{{ dag.dag_id }}', task_id='task-id', dag=self.dag, ) ti = TaskInstance(op, DEFAULT_DATE) ti.render_templates() self.assertEqual(dag_id, getattr(op, 'filter')['job_names'][0]) self.assertEqual(dag_id, getattr(op, 'gcp_conn_id'))
def test_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 = GcpTransferServiceOperationsListOperator( request_filter={"job_names": ['{{ dag.dag_id }}']}, gcp_conn_id='{{ dag.dag_id }}', task_id='task-id', dag=self.dag, ) ti = TaskInstance(op, DEFAULT_DATE) ti.render_templates() # pylint:disable=unsubscriptable-object self.assertEqual(dag_id, getattr(op, 'filter')['job_names'][0]) # pylint:enable=unsubscriptable-object self.assertEqual(dag_id, getattr(op, 'gcp_conn_id'))
# [START howto_operator_gcp_transfer_update_job] update_job = GcpTransferServiceJobUpdateOperator( task_id="update_job", job_name= "{{task_instance.xcom_pull('create_transfer_job_from_aws')['name']}}", body=update_body, ) # [END howto_operator_gcp_transfer_update_job] # [START howto_operator_gcp_transfer_list_operations] list_operations = GcpTransferServiceOperationsListOperator( task_id="list_operations", request_filter={ FILTER_PROJECT_ID: GCP_PROJECT_ID, FILTER_JOB_NAMES: [ "{{task_instance.xcom_pull('create_transfer_job_from_aws')['name']}}" ], }, ) # [END howto_operator_gcp_transfer_list_operations] # [START howto_operator_gcp_transfer_get_operation] get_operation = GcpTransferServiceOperationGetOperator( task_id="get_operation", operation_name= "{{task_instance.xcom_pull('list_operations')[0]['name']}}") # [END howto_operator_gcp_transfer_get_operation] # [START howto_operator_gcp_transfer_resume_operation] resume_operation = GcpTransferServiceOperationResumeOperator(