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 = CloudDataTransferServiceS3ToGCSOperator(
s3_bucket='{{ dag.dag_id }}',
gcs_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, 's3_bucket')
assert dag_id == getattr(op, 'gcs_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 = CloudDataTransferServiceS3ToGCSOperator(
task_id=TASK_ID,
s3_bucket=AWS_BUCKET_NAME,
gcs_bucket=GCS_BUCKET_NAME,
project_id=GCP_PROJECT_ID,
description=DESCRIPTION,
schedule=SCHEDULE_DICT,
)
assert operator.task_id == TASK_ID
assert operator.s3_bucket == AWS_BUCKET_NAME
assert operator.gcs_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 = CloudDataTransferServiceS3ToGCSOperator(
task_id=TASK_ID,
s3_bucket=AWS_BUCKET_NAME,
gcs_bucket=GCS_BUCKET_NAME,
project_id=GCP_PROJECT_ID,
description=DESCRIPTION,
schedule=SCHEDULE_DICT,
)
self.assertEqual(operator.task_id, TASK_ID)
self.assertEqual(operator.s3_bucket, AWS_BUCKET_NAME)
self.assertEqual(operator.gcs_bucket, GCS_BUCKET_NAME)
self.assertEqual(operator.project_id, GCP_PROJECT_ID)
self.assertEqual(operator.description, DESCRIPTION)
self.assertEqual(operator.schedule, SCHEDULE_DICT)
def test_execute(self, mock_aws_hook, mock_transfer_hook):
mock_aws_hook.return_value.get_credentials.return_value = Credentials(
TEST_AWS_ACCESS_KEY_ID, TEST_AWS_ACCESS_SECRET, None)
operator = CloudDataTransferServiceS3ToGCSOperator(
task_id=TASK_ID,
s3_bucket=AWS_BUCKET_NAME,
gcs_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_AWS_RAW)
assert mock_transfer_hook.return_value.wait_for_transfer_job.called
def test_execute_should_throw_ex_when_delete_job_without_wait(
self, mock_transfer_hook):
with pytest.raises(AirflowException) as ctx:
operator = CloudDataTransferServiceS3ToGCSOperator(
task_id=TASK_ID,
s3_bucket=AWS_BUCKET_NAME,
gcs_bucket=GCS_BUCKET_NAME,
description=DESCRIPTION,
schedule=SCHEDULE_DICT,
wait=False,
delete_job_after_completion=True,
)
operator.execute(None)
err = ctx.value
assert "If 'delete_job_after_completion' is True, then 'wait' must also be True." in str(
err)
mock_transfer_hook.assert_not_called()
