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()