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