def test_execute(self, mock_transfer_hook):
operator = GoogleCloudStorageToGoogleCloudStorageTransferOperator(
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
)
self.assertTrue(mock_transfer_hook.return_value.wait_for_transfer_job.called)
def test_execute(self, mock_transfer_hook):
operator = GoogleCloudStorageToGoogleCloudStorageTransferOperator(
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
)
self.assertTrue(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 = GoogleCloudStorageToGoogleCloudStorageTransferOperator(
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()
self.assertEqual(dag_id, getattr(op, 'source_bucket'))
self.assertEqual(dag_id, getattr(op, 'destination_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 = GoogleCloudStorageToGoogleCloudStorageTransferOperator(
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)
def test_templates(self, _):
dag_id = 'test_dag_id'
configuration.load_test_config()
args = {'start_date': DEFAULT_DATE}
self.dag = DAG(dag_id, default_args=args)
op = GoogleCloudStorageToGoogleCloudStorageTransferOperator(
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()
self.assertEqual(dag_id, getattr(op, 'source_bucket'))
self.assertEqual(dag_id, getattr(op, 'destination_bucket'))
self.assertEqual(dag_id, getattr(op, 'description'))
self.assertEqual(dag_id, getattr(op, 'object_conditions')['exclude_prefixes'][0])
self.assertEqual(dag_id, getattr(op, 'gcp_conn_id'))
