def test_s3_to_sftp_operation(self):
# Setting
configuration.conf.set("core", "enable_xcom_pickling", "True")
test_remote_file_content = \
"This is remote file content \n which is also multiline " \
"another line here \n this is last line. EOF"
# Test for creation of s3 bucket
conn = boto3.client('s3')
conn.create_bucket(Bucket=self.s3_bucket)
self.assertTrue((self.s3_hook.check_for_bucket(self.s3_bucket)))
with open(LOCAL_FILE_PATH, 'w') as file:
file.write(test_remote_file_content)
self.s3_hook.load_file(LOCAL_FILE_PATH,
self.s3_key,
bucket_name=BUCKET)
# Check if object was created in s3
objects_in_dest_bucket = conn.list_objects(Bucket=self.s3_bucket,
Prefix=self.s3_key)
# there should be object found, and there should only be one object found
self.assertEqual(len(objects_in_dest_bucket['Contents']), 1)
# the object found should be consistent with dest_key specified earlier
self.assertEqual(objects_in_dest_bucket['Contents'][0]['Key'],
self.s3_key)
# get remote file to local
run_task = S3ToSFTPOperator(s3_bucket=BUCKET,
s3_key=S3_KEY,
sftp_path=SFTP_PATH,
sftp_conn_id=SFTP_CONN_ID,
s3_conn_id=S3_CONN_ID,
task_id=TASK_ID,
dag=self.dag)
self.assertIsNotNone(run_task)
run_task.execute(None)
# Check that the file is created remotely
check_file_task = SSHOperator(task_id="test_check_file",
ssh_hook=self.hook,
command="cat {0}".format(self.sftp_path),
do_xcom_push=True,
dag=self.dag)
self.assertIsNotNone(check_file_task)
ti3 = TaskInstance(task=check_file_task,
execution_date=timezone.utcnow())
ti3.run()
self.assertEqual(
ti3.xcom_pull(task_ids='test_check_file',
key='return_value').strip(),
test_remote_file_content.encode('utf-8'))
# Clean up after finishing with test
conn.delete_object(Bucket=self.s3_bucket, Key=self.s3_key)
conn.delete_bucket(Bucket=self.s3_bucket)
self.assertFalse((self.s3_hook.check_for_bucket(self.s3_bucket)))
"AWS_SECRET_ACCESS_KEY": AIRFLOW_S3.password,
"DATETIME": "{{ ti.xcom_pull(task_ids='set_datetime') }}"
}
},
dag=DAG)
# TODO: Replace with S3ListOperator
NOTE_TYPES = ["collection", "service"]
S3_TO_SERVER_TASKS = {}
for note_type in NOTE_TYPES:
S3_TO_SERVER_TASKS[note_type] = S3ToSFTPOperator(
task_id='s3_to_server_' + note_type,
sftp_conn_id="tul_cob",
sftp_path="/tmp/" + note_type + "_notes.json",
s3_conn_id="AIRFLOW_S3",
s3_bucket=AIRFLOW_DATA_BUCKET,
s3_key="electronic-notes/{{ ti.xcom_pull(task_ids='set_datetime') }}/"
+ note_type + "_notes.json",
dag=DAG)
S3_TO_SERVER_TASKS = list(S3_TO_SERVER_TASKS.values())
RELOAD_ELECTRONIC_COLLECTIONS_COMMAND = """
sudo su - tul_cob bash -c \
"cd /var/www/tul_cob && RAILS_ENV=production bundle exec rake reload_electronic_notes"
"""
RELOAD_ELECTRONIC_NOTES = SSHOperator(
task_id="reload_electronic_notes",
command=RELOAD_ELECTRONIC_COLLECTIONS_COMMAND,