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)