def test_execute(self, mock_transfer_hook):
operator = CloudDataTransferServiceGCSToGCSOperator(
task_id=TASK_ID,
source_bucket=GCS_BUCKET_NAME,
destination_bucket=GCS_BUCKET_NAME,
description=DESCRIPTION,
schedule=SCHEDULE_DICT,
)
operator.execute(None)
mock_transfer_hook.return_value.create_transfer_job.assert_called_once_with(
body=VALID_TRANSFER_JOB_GCS_RAW)
assert mock_transfer_hook.return_value.wait_for_transfer_job.called
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 = CloudDataTransferServiceGCSToGCSOperator(
source_bucket='{{ dag.dag_id }}',
destination_bucket='{{ dag.dag_id }}',
description='{{ dag.dag_id }}',
object_conditions={'exclude_prefixes': ['{{ dag.dag_id }}']},
gcp_conn_id='{{ dag.dag_id }}',
task_id=TASK_ID,
dag=self.dag,
)
ti = TaskInstance(op, DEFAULT_DATE)
ti.render_templates()
assert dag_id == getattr(op, 'source_bucket')
assert dag_id == getattr(op, 'destination_bucket')
assert dag_id == getattr(op, 'description')
# pylint: disable=unsubscriptable-object
assert dag_id == getattr(op,
'object_conditions')['exclude_prefixes'][0]
# pylint: enable=unsubscriptable-object
assert dag_id == getattr(op, 'gcp_conn_id')
def test_constructor(self):
operator = CloudDataTransferServiceGCSToGCSOperator(
task_id=TASK_ID,
source_bucket=GCS_BUCKET_NAME,
destination_bucket=GCS_BUCKET_NAME,
project_id=GCP_PROJECT_ID,
description=DESCRIPTION,
schedule=SCHEDULE_DICT,
)
assert operator.task_id == TASK_ID
assert operator.source_bucket == GCS_BUCKET_NAME
assert operator.destination_bucket == GCS_BUCKET_NAME
assert operator.project_id == GCP_PROJECT_ID
assert operator.description == DESCRIPTION
assert operator.schedule == SCHEDULE_DICT
def test_constructor(self):
operator = CloudDataTransferServiceGCSToGCSOperator(
task_id=TASK_ID,
source_bucket=GCS_BUCKET_NAME,
destination_bucket=GCS_BUCKET_NAME,
project_id=GCP_PROJECT_ID,
description=DESCRIPTION,
schedule=SCHEDULE_DICT,
)
self.assertEqual(operator.task_id, TASK_ID)
self.assertEqual(operator.source_bucket, GCS_BUCKET_NAME)
self.assertEqual(operator.destination_bucket, GCS_BUCKET_NAME)
self.assertEqual(operator.project_id, GCP_PROJECT_ID)
self.assertEqual(operator.description, DESCRIPTION)
self.assertEqual(operator.schedule, SCHEDULE_DICT)
