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 = S3ToGoogleCloudStorageTransferOperator(
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()
self.assertEqual(dag_id, getattr(op, 's3_bucket'))
self.assertEqual(dag_id, getattr(op, 'gcs_bucket'))
self.assertEqual(dag_id, getattr(op, 'description'))
# pylint:disable=unsubscriptable-object
self.assertEqual(
dag_id,
getattr(op, 'object_conditions')['exclude_prefixes'][0])
# pylint:enable=unsubscriptable-object
self.assertEqual(dag_id, getattr(op, 'gcp_conn_id'))
def test_constructor(self):
operator = S3ToGoogleCloudStorageTransferOperator(
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 = S3ToGoogleCloudStorageTransferOperator(
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)
self.assertTrue(
mock_transfer_hook.return_value.wait_for_transfer_job.called)
